Skip to content

New 2d tracking #237

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
leoscholl wants to merge 23 commits into
base: master
Choose a base branch
from
Open

New 2d tracking #237

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
d921338
generate trajectories in two dimensions
mtringi Jul 21, 2025
01f4d38
add black and invisible? update target to move in two directions
leoscholl Jul 21, 2025
c7f6cc0
change original function to match master
mtringi Jul 26, 2025
6d04685
Merge branch 'tracking_2D' of github.com:aolabNeuro/brain-python-inte…
mtringi Jul 26, 2025
653cbbe
fix lookahead
leoscholl Jul 29, 2025
5de3690
cleanup
leoscholl Aug 1, 2025
0459c03
snake always in front
leoscholl Aug 4, 2025
7116526
add logic to distribute frequencies
mtringi Aug 12, 2025
7eb11f5
Merge branch 'tracking_2D' of github.com:aolabNeuro/brain-python-inte…
mtringi Aug 12, 2025
c2f7b50
graphics fixes
leoscholl Aug 12, 2025
a9a6a1f
add fft plots to test_tasks
mtringi Aug 12, 2025
697b84e
marios fix
leoscholl Aug 14, 2025
bebc2f0
try exponential decay for amplitude of sine waves
mtringi Oct 10, 2025
9abeb41
Merge branch 'master' into tracking_2D
leoscholl Nov 17, 2025
e969887
fix hideleft feature
leoscholl Nov 17, 2025
6af1654
update 2d tracking to include optional exponential amplitudes and use…
mtringi Dec 6, 2025
65159ca
fix lookahead
leoscholl Dec 6, 2025
8cfb6f9
fix frame shift in traj bug, added user_screen to task data, saved te…
katherineperks Dec 17, 2025
32aeca2
fix transition test funcs, save test data with ramps
katherineperks Dec 20, 2025
17b4e10
Merge remote-tracking branch 'origin/master' into tracking_2D to incl…
mtringi Jan 9, 2026
b481c7e
increment bmi3d version # and comment out print statements
katherineperks Jan 12, 2026
aa17751
Merge branch 'tracking_2D' into test-merge-tracking-task
katherineperks Jan 12, 2026
64b2aca
update trajectory names and scale
katherineperks Jan 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions built_in_tasks/manualcontrolmultitasks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ def __init__(self, *args, **kwargs):

def init(self):
self.add_dtype('manual_input', 'f8', (3,))
self.add_dtype('user_screen', 'f8', (3,))
super().init()
self.no_data_counter = np.zeros((self._quality_window_size,), dtype='?')

Expand Down Expand Up @@ -121,13 +122,15 @@ def move_effector(self, pos_offset=[0,0,0], vel_offset=[0,0,0]):
self.no_data_counter[self.cycle_count % self._quality_window_size] = 1
self.update_report_stats()
self.task_data['manual_input'] = np.ones((3,))*np.nan
self.task_data['user_screen'] = np.ones((3,))*np.nan
return

self.task_data['manual_input'] = raw_coords.copy()
self.no_data_counter[self.cycle_count % self._quality_window_size] = 0

# Transform coordinates
coords = self._transform_coords(raw_coords)
self.task_data['user_screen'] = coords

try:
if self.limit2d:
Expand Down
20 changes: 18 additions & 2 deletions built_in_tasks/target_graphics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
Base tasks for generic point-to-point reaching
'''
import numpy as np
from riglib.stereo_opengl.primitives import Cable, Sphere, Cube, Torus, Text
from riglib.stereo_opengl.primitives import Cable, Snake, Sphere, Cube, Torus, Text
from riglib.stereo_opengl.primitives import Cylinder, Plane, Sphere, Cube
from riglib.stereo_opengl.models import FlatMesh, Group
from riglib.stereo_opengl.textures import Texture, TexModel
Expand Down Expand Up @@ -36,6 +36,8 @@
"gold": (0.941,0.637,0.25,0.75),
"elephant":(0.5,0.5,0.5,0.5),
"white": (1, 1, 1, 0.75),
"black": (0, 0, 0, 0.75),
"invisible": (0, 0, 0, 0.0),
}

class CircularTarget(object):
Expand Down Expand Up @@ -186,7 +188,7 @@ def __init__(self, target_radius=1, target_color=(1, 0, 0, .5), starting_pos=np.
self._pickle_init()

def _pickle_init(self):
self.cable = Cable(radius=self.target_radius,trajectory = self.trajectory, color=self.target_color)
self.cable = Cable(radius=self.target_radius, xyz=self.trajectory, color=self.target_color)
self.graphics_models = [self.cable]
self.cable.translate(*self.position)

Expand Down Expand Up @@ -238,6 +240,20 @@ def reset(self):
def get_position(self):
return self.cable.xfm.move

class VirtualSnakeTarget(VirtualCableTarget):

def _pickle_init(self):
self.trajectory = np.array(self.trajectory)
self.cable = Snake(radius=self.target_radius, trajectory=self.trajectory, color=self.target_color)
self.graphics_models = [self.cable]
self.cable.translate(*self.position)

def update_mask(self, start_frame, end_frame, inverse=False):
'''
Update the texture mask of the snake target.
'''
self.cable.update_texture(start_frame, end_frame, inverse=inverse)

class VirtualTorusTarget(VirtualCircularTarget):

def __init__(self, inner_radius=2, outer_radius=3, target_color=(1, 0, 0, .5), starting_pos=np.zeros(3)):
Expand Down
343 changes: 243 additions & 100 deletions built_in_tasks/target_tracking_task.py
View file
Open in desktop

Large diffs are not rendered by default.

14 changes: 9 additions & 5 deletions features/generator_features.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -328,15 +328,19 @@ def _start_reward(self):

class HideLeftTrajectory(traits.HasTraits):
'''
Cover left side of tracking task screen with a black box.
Hide the left side of the tracking trajectory.
This will cover the 'lookbehind' of the target trajectory.
Useful for task with bumpers.
'''

def _start_trajectory(self):
super()._start_trajectory()
if self.frame_index == 0:
self.box.show()
def setup_start_wait(self):
super().setup_start_wait()
print(self.frame_index)
self.trajectory.update_mask(self.lookahead+2, self.lookahead*2)

def update_frame(self):
super().update_frame()
self.trajectory.update_mask(self.frame_index+self.lookahead+1, self.frame_index+2*self.lookahead)

class ReadysetMedley(traits.HasTraits):

Expand Down
16 changes: 13 additions & 3 deletions riglib/stereo_opengl/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,9 @@ def __init__(self, shader="default", color=(0.5, 0.5, 0.5, 1),
# The orientation of the object, in the world frame
self._xfm = self.xfm
self.allocated = False

# Keep track of the model's size for rendering
self.bounding_radius = 0.0

def __setattr__(self, attr, xfm):
'''Checks if the xfm was changed, and recaches the _xfm which is sent to the shader'''
Expand Down Expand Up @@ -185,13 +188,18 @@ def init(self):
model.init()

def render_queue(self, xfm=np.eye(4), **kwargs):
for model in self.models:
def sort_key(obj):
pos = obj.xfm.move[1]
radius = obj.bounding_radius
dist = pos - radius
return -dist # Negative for far-to-near sorting
sorted_models = sorted(self.models, key=sort_key)
for model in sorted_models:
for out in model.render_queue(**kwargs):
yield out

def draw(self, ctx, **kwargs):
sorted_models = sorted(self.models, key=lambda obj: -obj.xfm.move[1])
for model in sorted_models:
for model in self.models:
model.draw(ctx, **kwargs)

def __getitem__(self, idx):
Expand Down Expand Up @@ -232,6 +240,8 @@ def __init__(self, verts, polys, normals=None, tcoords=None, **kwargs):
self.polys = polys
self.tcoords = tcoords
self.normals = normals

self.bounding_radius = abs(np.min(self.verts[:, 1]))

def init(self):
allocated = super(TriMesh, self).init()
Expand Down
119 changes: 97 additions & 22 deletions riglib/stereo_opengl/primitives.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@

from .models import TriMesh
from .textures import Texture, TexModel
from OpenGL.GL import GL_NEAREST, GL_REPEAT
from OpenGL.GL import *
from PIL import Image, ImageDraw, ImageFont
import matplotlib.font_manager as fm

Expand Down Expand Up @@ -154,33 +154,68 @@ def __init__(self, height=1, radius=1, segments=36, **kwargs):
super().__init__(total_pts, total_polys, tcoords=total_tcoords, normals=total_normals, **kwargs)

class Cable(TriMesh):
def __init__(self,radius=.5, trajectory = np.array([np.sin(x) for x in range(60)]), segments=12,**kwargs):
self.trial_trajectory = trajectory
def __init__(self,radius=.5, xyz = np.array([np.sin(x) for x in range(60)]), segments=12,**kwargs):
self.xyz = xyz
if np.ndim(xyz) == 1:
self.xyz = np.array([[x,0,xyz[x]] for x in range(len(xyz))])
self.center_value = [0,0,0]
self.radius = radius
self.segments = segments
self.update(**kwargs)

def update(self, **kwargs):
theta = np.linspace(0, 2*np.pi, self.segments, endpoint=False)
unit = np.array([np.ones(self.segments),np.cos(theta) ,np.sin(theta)]).T
intial = np.array([[0,0,self.trial_trajectory[x]] for x in range(len(self.trial_trajectory))])
self.pts = (unit*[-30/1.36,self.radius,self.radius])+intial[0]
for i in range(1,len(intial)):
self.pts = np.vstack([self.pts, (unit*[(i-30)/3,self.radius,self.radius])+intial[i]])

self.normals = np.vstack([unit*[1,1,0], unit*[1,1,0]])
self.polys = []
for i in range(self.segments-1):
for j in range(len(intial)-1):
self.polys.append((i+j*self.segments, i+1+j*self.segments, i+self.segments+j*self.segments))
self.polys.append((i+self.segments+j*self.segments, i+1+j*self.segments, i+1+self.segments+j*self.segments))

tcoord = np.array([np.arange(self.segments), np.ones(self.segments)]).T
n = 1./self.segments
self.tcoord = np.vstack([tcoord*[n,1], tcoord*[n,0]])
super(Cable, self).__init__(self.pts, np.array(self.polys),
tcoords=self.tcoord, normals=self.normals, **kwargs)
theta = np.linspace(0, 2 * np.pi, self.segments, endpoint=False)
circle = np.stack([np.cos(theta), np.sin(theta)], axis=1) # (segments, 2)

pts = []
normals = []
tcoords = []
n_path = len(self.xyz)

a = np.array([0, 1, 0]) # fixed up direction

# Compute tangents along path
tangents = np.gradient(self.xyz, axis=0)
tangents = tangents / np.linalg.norm(tangents, axis=1, keepdims=True)

for i in range(n_path):
p = self.xyz[i]
t = tangents[i]

# Ring orientation
b = np.cross(t, a)
if np.linalg.norm(b) < 1e-6:
b = np.array([0, 0, 1]) # fallback
else:
b = b / np.linalg.norm(b)

for j in range(self.segments):
cx, cy = circle[j]
offset = self.radius * (cx * b + cy * a)
pts.append(p + offset)
normals.append(offset / np.linalg.norm(offset))
tcoords.append([j / self.segments, i / (n_path - 1)])

self.pts = np.array(pts)
self.normals = np.array(normals)
self.tcoord = np.array(tcoords)

# Create triangle strips between rings
polys = []
for i in range(n_path - 1):
for j in range(self.segments):
i0 = i * self.segments + j
i1 = i * self.segments + (j + 1) % self.segments
i2 = (i + 1) * self.segments + j
i3 = (i + 1) * self.segments + (j + 1) % self.segments

polys.append((i2, i1, i0))
polys.append((i3, i1, i2))

self.polys = np.array(polys)

super().__init__(self.pts, self.polys, tcoords=self.tcoord,
normals=self.normals, **kwargs)

class Torus(TriMesh):
'''
Expand Down Expand Up @@ -473,6 +508,46 @@ def __init__(self, radius, alpha=1, stop=False, **kwargs):
texture_mapping='planar', **kwargs)
self.rotate_x(90) # Make the target face the camera

class Snake(Cable, TexModel):
'''
A Cable with a gradient texture applied along its length.
'''
def __init__(self, radius=.5, trajectory=np.array([np.sin(x) for x in range(100)]), segments=12, **kwargs):
self.trajectory = trajectory
color = kwargs.pop('color', [1, 1, 1, 1]) # Default color if not provided
self.color = color
tex = self.get_texture(0, len(trajectory))
super().__init__(radius, trajectory, segments, tex=tex, color=[0, 0, 0, 1], **kwargs)
self.color = color # Store the color for later use

def get_texture(self, start_frame, end_frame, inverse=False):
mask = np.zeros((len(self.trajectory)))
if start_frame >= len(self.trajectory):
start_frame = len(self.trajectory)
if end_frame >= len(self.trajectory):
end_frame = len(self.trajectory)
mask[start_frame:end_frame] = 1
if inverse:
mask = 1 - mask
mask = np.tile(mask, (4, 1)).T # Repeat for RGBA
mask = self.color * mask # Apply color
tex = Texture(mask.reshape((1, len(mask), 4))) # Reshape to (1, n_colors, 4)
return tex

def update_texture(self, start_frame, end_frame, inverse=False):
'''
Update the texture of the snake based on the new trajectory.
'''
self.tex.delete() # Delete the old texture
tex = self.get_texture(start_frame, end_frame, inverse=inverse)
self.tex = tex
self.tex.init()

def draw(self, ctx):
glDisable(GL_DEPTH_TEST)
super().draw(ctx)
glEnable(GL_DEPTH_TEST)

##### 2-D primitives #####

class Shape2D(object):
Expand Down
2 changes: 2 additions & 0 deletions riglib/stereo_opengl/shaders/none.f.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,5 +13,7 @@ void main() {
texcolor.rgb + basecolor.rgb,
texcolor.a * basecolor.a
);
if (frag_diffuse.a < 0.01)
discard;
FragColor = frag_diffuse;
}
2 changes: 2 additions & 0 deletions riglib/stereo_opengl/shaders/phong.f.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,8 @@ vec4 phong() {
texcolor.rgb + basecolor.rgb,
texcolor.a * basecolor.a
);
if (frag_diffuse.a < 0.01)
discard;

vec4 diffuse_factor
= max(-dot(normal, mv_light_direction), 0.0) * light_diffuse;
Expand Down
12 changes: 7 additions & 5 deletions riglib/stereo_opengl/textures.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,11 +26,13 @@ def __init__(self, tex, size=None,

if isinstance(tex, np.ndarray):
if tex.max() <= 1:
tex *= 255
if len(tex.shape) < 3:
tex = np.tile(tex, [3, 1, 1]).T
if tex.shape[-1] == 3:
tex = np.dstack([tex, np.ones(tex.shape[:-1])])
tex = (tex * 255).astype(np.uint8)
else:
tex = tex.astype(np.uint8)
if tex.ndim == 2:
tex = np.stack([tex]*3, axis=-1) # grayscale → RGB
elif tex.shape[-1] == 1:
tex = np.repeat(tex, 3, axis=-1)
size = tex.shape[:2]
tex = tex.astype(np.uint8).tobytes()
elif isinstance(tex, str):
Expand Down
2 changes: 1 addition & 1 deletion riglib/stereo_opengl/window.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ def screen_init(self):

glDisable(GL_FRAMEBUFFER_SRGB) # disable gamma correction
glEnable(GL_BLEND)
glDepthFunc(GL_LESS)
glDepthFunc(GL_LEQUAL)
glEnable(GL_DEPTH_TEST)
glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
glClearColor(*self.background)
Expand Down
2 changes: 1 addition & 1 deletion setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

setuptools.setup(
name="aolab-bmi3d",
version="1.2.4",
version="1.2.5",
author="Lots of people",
description="electrophysiology experimental rig library",
packages=setuptools.find_packages(),
Expand Down
18 changes: 12 additions & 6 deletions tests/test_graphics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@

from riglib.stereo_opengl.environment import Grid
from riglib.stereo_opengl.window import Window, Window2D, FPScontrol
from riglib.stereo_opengl.primitives import AprilTag, Cylinder, Cube, Plane, Sphere, Cone, Text, TexSphere, TexCube, TexPlane
from riglib.stereo_opengl.primitives import AprilTag, Cylinder, Cube, Plane, Snake, Sphere, Cone, Text, Cable, TexSphere, TexCube, TexPlane
from features.optitrack_features import SpheresToCylinders
from riglib.stereo_opengl.window import Window, Window2D, FPScontrol, WindowSSAO
from riglib.stereo_opengl.openxr import WindowVR
Expand Down Expand Up @@ -39,13 +39,14 @@
planet = Sphere(3, color=[0.75,0.25,0.25,0.75])
orbit_radius = 4
orbit_speed = 1
wobble_radius = 0
wobble_radius = 5
wobble_speed = 0.5
#TexSphere = type("TexSphere", (Sphere, TexModel), {})
#TexPlane = type("TexPlane", (Plane, TexModel), {})
#reward_text = Text(7.5, "123", justify='right', color=[1,0,1,1])
reward_text = Text(7.5, "123", justify='right', color=[1,0,1,1])
# center_out_gen = ScreenTargetCapture.centerout_2D(1)
# center_out_positions = [pos[1] for _, pos in center_out_gen]
cable = Snake(0.5, 2*np.sin(np.arange(200)/2), color=(1,0,1,0.75)).translate(-15, 0, -10)
center_out_gen = ScreenTargetCapture.centerout_tabletop(1)
center_out_positions = [(pos[1][0], pos[1][1], -10) for _, pos in center_out_gen]
center_out_targets = [
Expand All @@ -70,8 +71,9 @@ def _start_draw(self):
#arm4j.set_joint_pos([0,0,np.pi/2,np.pi/2])
#arm4j.get_endpoint_pos()
self.add_model(Grid(50))
self.add_model(apriltag)
self.add_model(moon)
self.add_model(planet)
self.add_model(apriltag)
# self.add_model(moon)
# self.add_model(planet)
# self.add_model(arm4j)
Expand All @@ -80,16 +82,20 @@ def _start_draw(self):
# self.add_model(TexPlane(5,5, tex=cloudy_tex(), specular_color=(0.,0,0,1)).rotate_x(90))
# self.add_model(TexPlane(5,5, specular_color=(0.,0,0,1), tex=cloudy_tex()).rotate_x(90))
# reward_text = Text(7.5, "123", justify='right', color=[1,0,1,0.75])
# self.add_model(reward_text)
self.add_model(reward_text)
# self.add_model(TexPlane(4,4,color=[0,0,0,0.9], tex=cloudy_tex()).rotate_x(90).translate(0,0,-5))
#self.screen_init()
#self.draw_world()
for model in center_out_targets:
self.add_model(model)
self.add_model(Sphere(radius=1, color=target_colors['purple']).translate(3,3,-10))
for model in center_out_targets:
self.add_model(model)
self.add_model(Sphere(radius=1, color=target_colors['purple']).translate(3,0,-10))
self.add_model(cable)

def _while_draw(self):
ts = time.time() - self.start_time
# cable.update_texture(int(ts*10), len(cable.trajectory))

x = travel_radius * np.cos(ts * travel_speed)
y = travel_radius * np.sin(ts * travel_speed)
Expand Down
Loading