Skip to content

Implemented environment-aware feature for Categorizer #336

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

Merged
yujiahu415 merged 10 commits into from
Jan 9, 2026
Merged

Implemented environment-aware feature for Categorizer #336

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
f00ce6b
Update categorizer.py
yujiahu415 Jan 8, 2026
9a5722d
Update categorizer.py
yujiahu415 Jan 8, 2026
0ecdec4
Update categorizer.py
yujiahu415 Jan 8, 2026
e43d081
Update categorizer.py
yujiahu415 Jan 8, 2026
be3bc37
Update gui_categorizer.py
yujiahu415 Jan 8, 2026
cab9c81
Update gui_categorizer.py
yujiahu415 Jan 8, 2026
e2f6de1
Update analyzebehavior_dt.py
yujiahu415 Jan 8, 2026
7afab80
Update analyzebehavior_dt.py
yujiahu415 Jan 8, 2026
62bc857
Update gui_analyzer.py and gui_categorizer.py
yujiahu415 Jan 8, 2026
b5c78d4
Update gui_analyzer.py
yujiahu415 Jan 8, 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
21 changes: 15 additions & 6 deletions LabGym/analyzebehavior_dt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -492,13 +492,14 @@ def detect_track_individuals(self,frames,batch_size,frame_count_analyze,backgrou
self.animations[animal_name][i][frame_count_analyze+1-batch_size+batch_count]=np.array(animation)


def detect_track_interact(self,frames,batch_size,frame_count_analyze,background_free=True,black_background=True):
def detect_track_interact(self,frames,batch_size,frame_count_analyze,background_free=True,black_background=True,color_costar=False):

# frames: frames that the Detector runs on
# batch_size: for batch inferencing by the Detector
# frame_count_analyze: the analyzed frame count
# background_free: whether to include background in animations
# black_background: whether to set background black
# color_costar: in 'interactive advanced' mode, whether to make the supporting roles RGB scale in animations

tensor_frames=[torch.as_tensor(frame.astype('float32').transpose(2,0,1)) for frame in frames]
inputs=[{'image':tensor_frame} for tensor_frame in tensor_frames]
Expand Down Expand Up @@ -626,7 +627,10 @@ def detect_track_interact(self,frames,batch_size,frame_count_analyze,background_
if background_free:
blob=frame*cv2.cvtColor(all_masks[i],cv2.COLOR_GRAY2BGR)
if other_mask is not None:
other_blob=cv2.cvtColor(cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)*other_mask,cv2.COLOR_GRAY2BGR)
if color_costar:
other_blob=frame*cv2.cvtColor(other_mask,cv2.COLOR_GRAY2BGR)
else:
other_blob=cv2.cvtColor(cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)*other_mask,cv2.COLOR_GRAY2BGR)
blob=cv2.add(blob,other_blob)
if black_background is False:
if other_mask is not None:
Expand Down Expand Up @@ -677,11 +681,12 @@ def detect_track_interact(self,frames,batch_size,frame_count_analyze,background_
self.track_animal_interact(frame_count_analyze+1-batch_size+batch_count,all_contours,other_contours,all_centers,all_heights,inners=all_inners,other_inners=other_inners,blobs=all_blobs)


def acquire_information(self,batch_size=1,background_free=True,black_background=True):
def acquire_information(self,batch_size=1,background_free=True,black_background=True,color_costar=False):

# batch_size: for batch inferencing by the Detector
# background_free: whether to include background in animations
# black_background: whether to set background black
# color_costar: in 'interactive advanced' mode, whether to make the supporting roles RGB scale in animations

print('Acquiring information in each frame...')
self.log.append('Acquiring information in each frame...')
Expand Down Expand Up @@ -728,7 +733,7 @@ def acquire_information(self,batch_size=1,background_free=True,black_background=
if batch_count==batch_size:
batch_count=0
if self.behavior_mode==2:
self.detect_track_interact(batch,batch_size,frame_count_analyze,background_free=background_free,black_background=black_background)
self.detect_track_interact(batch,batch_size,frame_count_analyze,background_free=background_free,black_background=black_background,color_costar=color_costar)
else:
self.detect_track_individuals(batch,batch_size,frame_count_analyze,background_free=background_free,black_background=black_background,animation=animation)
batch=[]
Expand Down Expand Up @@ -1915,11 +1920,12 @@ def generate_data_interact_basic(self,background_free=True,black_background=True
print('Behavior example generation completed!')


def generate_data_interact_advance(self,background_free=True,black_background=True,skip_redundant=1):
def generate_data_interact_advance(self,background_free=True,black_background=True,skip_redundant=1,color_costar=False):

# background_free: whether to include background in animations
# black_background: whether to set background black
# skip_redundant: the interval (in frames) of two consecutively generated behavior example pairs
# color_costar: in 'interactive advanced' mode, whether to make the supporting roles RGB scale in animations

print('Generating behavior examples...')
print(datetime.datetime.now())
Expand Down Expand Up @@ -2073,7 +2079,10 @@ def generate_data_interact_advance(self,background_free=True,black_background=Tr
if background_free:
blob=frame*cv2.cvtColor(all_masks[i],cv2.COLOR_GRAY2BGR)
if other_mask is not None:
other_blob=cv2.cvtColor(cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)*other_mask,cv2.COLOR_GRAY2BGR)
if color_costar:
other_blob=frame*cv2.cvtColor(other_mask,cv2.COLOR_GRAY2BGR)
else:
other_blob=cv2.cvtColor(cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)*other_mask,cv2.COLOR_GRAY2BGR)
blob=cv2.add(blob,other_blob)
if black_background is False:
if other_mask is not None:
Expand Down
48 changes: 40 additions & 8 deletions LabGym/categorizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1193,7 +1193,7 @@ def train_pattern_recognizer(self,data_path,model_path,out_path=None,dim=64,chan
self.train_pattern_recognizer_onfly(out_folder,model_path,out_path=out_path,dim=dim,channel=channel,time_step=time_step,level=level,include_bodyparts=include_bodyparts,std=std,background_free=background_free,black_background=black_background,behavior_mode=behavior_mode,social_distance=social_distance)


def train_animation_analyzer(self,data_path,model_path,out_path=None,dim=64,channel=1,time_step=15,level=2,aug_methods=[],augvalid=True,include_bodyparts=True,std=0,background_free=True,black_background=True,behavior_mode=0,social_distance=0,out_folder=None):
def train_animation_analyzer(self,data_path,model_path,out_path=None,dim=64,channel=1,time_step=15,level=2,aug_methods=[],augvalid=True,include_bodyparts=True,std=0,background_free=True,black_background=True,behavior_mode=0,social_distance=0,color_costar=False,out_folder=None):

# data_path: the folder that stores all the prepared training examples
# model_path: the path to the trained Animation Analyzer
Expand All @@ -1210,6 +1210,7 @@ def train_animation_analyzer(self,data_path,model_path,out_path=None,dim=64,chan
# black_background: whether to set background
# behavior_mode: 0--non-interactive, 1--interactive basic, 2--interactive advanced, 3--static images
# social_distance: a threshold (folds of size of a single animal) on whether to include individuals that are not main character in behavior examples
# color_costar: in 'interactive advanced' mode, whether to make the supporting roles RGB scale in animations
# out_folder: if not None, will output all the augmented data to this folder

filters=8
Expand Down Expand Up @@ -1264,7 +1265,12 @@ def train_animation_analyzer(self,data_path,model_path,out_path=None,dim=64,chan
else:
black_code=1

parameters={'classnames':list(self.classnames),'dim_tconv':int(dim),'channel':int(channel),'time_step':int(time_step),'network':1,'level_tconv':int(level),'inner_code':int(inner_code),'std':int(std),'background_free':int(background_code),'black_background':int(black_code),'behavior_kind':int(behavior_mode),'social_distance':int(social_distance)}
if color_costar:
color_code=0
else:
color_code=1

parameters={'classnames':list(self.classnames),'dim_tconv':int(dim),'channel':int(channel),'time_step':int(time_step),'network':1,'level_tconv':int(level),'inner_code':int(inner_code),'std':int(std),'background_free':int(background_code),'black_background':int(black_code),'behavior_kind':int(behavior_mode),'social_distance':int(social_distance),'color_code':int(color_code)}
pd_parameters=pd.DataFrame.from_dict(parameters)
pd_parameters.to_csv(os.path.join(model_path,'model_parameters.txt'),index=False)

Expand Down Expand Up @@ -1394,7 +1400,7 @@ def train_animation_analyzer(self,data_path,model_path,out_path=None,dim=64,chan
self.train_animation_analyzer_onfly(out_folder,model_path,out_path=out_path,dim=dim,channel=channel,time_step=time_step,level=level,include_bodyparts=include_bodyparts,std=std,background_free=background_free,black_background=black_background,behavior_mode=behavior_mode,social_distance=social_distance)


def train_combnet(self,data_path,model_path,out_path=None,dim_tconv=32,dim_conv=64,channel=1,time_step=15,level_tconv=1,level_conv=2,aug_methods=[],augvalid=True,include_bodyparts=True,std=0,background_free=True,black_background=True,behavior_mode=0,social_distance=0,out_folder=None):
def train_combnet(self,data_path,model_path,out_path=None,dim_tconv=32,dim_conv=64,channel=1,time_step=15,level_tconv=1,level_conv=2,aug_methods=[],augvalid=True,include_bodyparts=True,std=0,background_free=True,black_background=True,behavior_mode=0,social_distance=0,color_costar=False,out_folder=None):

# data_path: the folder that stores all the prepared training examples
# model_path: the path to the trained Categorizer
Expand All @@ -1413,6 +1419,7 @@ def train_combnet(self,data_path,model_path,out_path=None,dim_tconv=32,dim_conv=
# black_background: whether to set background black
# behavior_mode: 0--non-interactive, 1--interactive basic, 2--interactive advanced, 3--static images
# social_distance: a threshold (folds of size of a single animal) on whether to include individuals that are not main character in behavior examples
# color_costar: in 'interactive advanced' mode, whether to make the supporting roles RGB scale in animations
# out_folder: if not None, will output all the augmented data to this folder

print('Training Categorizer with both Animation Analyzer and Pattern Recognizer using the behavior examples in: '+str(data_path))
Expand Down Expand Up @@ -1460,7 +1467,12 @@ def train_combnet(self,data_path,model_path,out_path=None,dim_tconv=32,dim_conv=
else:
black_code=1

parameters={'classnames':list(self.classnames),'dim_tconv':int(dim_tconv),'dim_conv':int(dim_conv),'channel':int(channel),'time_step':int(time_step),'network':2,'level_tconv':int(level_tconv),'level_conv':int(level_conv),'inner_code':int(inner_code),'std':int(std),'background_free':int(background_code),'black_background':int(black_code),'behavior_kind':int(behavior_mode),'social_distance':int(social_distance)}
if color_costar:
color_code=0
else:
color_code=1

parameters={'classnames':list(self.classnames),'dim_tconv':int(dim_tconv),'dim_conv':int(dim_conv),'channel':int(channel),'time_step':int(time_step),'network':2,'level_tconv':int(level_tconv),'level_conv':int(level_conv),'inner_code':int(inner_code),'std':int(std),'background_free':int(background_code),'black_background':int(black_code),'behavior_kind':int(behavior_mode),'social_distance':int(social_distance),'color_code':int(color_code)}
pd_parameters=pd.DataFrame.from_dict(parameters)
pd_parameters.to_csv(os.path.join(model_path,'model_parameters.txt'),index=False)

Expand Down Expand Up @@ -1705,7 +1717,7 @@ def train_pattern_recognizer_onfly(self,data_path,model_path,out_path=None,dim=3
print('No train / validation folder!')


def train_animation_analyzer_onfly(self,data_path,model_path,out_path=None,dim=32,channel=1,time_step=15,level=2,include_bodyparts=True,std=0,background_free=True,black_background=True,behavior_mode=0,social_distance=0):
def train_animation_analyzer_onfly(self,data_path,model_path,out_path=None,dim=32,channel=1,time_step=15,level=2,include_bodyparts=True,std=0,background_free=True,black_background=True,behavior_mode=0,social_distance=0,color_costar=False):

# data_path: the folder that stores all the prepared training examples
# model_path: the path to the trained Animation Analyzer
Expand All @@ -1720,6 +1732,7 @@ def train_animation_analyzer_onfly(self,data_path,model_path,out_path=None,dim=3
# black_background: whether to set background
# behavior_mode: 0--non-interactive, 1--interactive basic, 2--interactive advanced, 3--static images
# social_distance: a threshold (folds of size of a single animal) on whether to include individuals that are not main character in behavior examples
# color_costar: in 'interactive advanced' mode, whether to make the supporting roles RGB scale in animations

filters=8

Expand Down Expand Up @@ -1765,7 +1778,12 @@ def train_animation_analyzer_onfly(self,data_path,model_path,out_path=None,dim=3
else:
black_code=1

parameters={'classnames':list(train_data.classmapping.keys()),'dim_tconv':int(dim),'channel':int(channel),'time_step':int(time_step),'network':1,'level_tconv':int(level),'inner_code':int(inner_code),'std':int(std),'background_free':int(background_code),'black_background':int(black_code),'behavior_kind':int(behavior_mode),'social_distance':int(social_distance)}
if color_costar:
color_code=0
else:
color_code=1

parameters={'classnames':list(train_data.classmapping.keys()),'dim_tconv':int(dim),'channel':int(channel),'time_step':int(time_step),'network':1,'level_tconv':int(level),'inner_code':int(inner_code),'std':int(std),'background_free':int(background_code),'black_background':int(black_code),'behavior_kind':int(behavior_mode),'social_distance':int(social_distance),'color_code':int(color_code)}
pd_parameters=pd.DataFrame.from_dict(parameters)
pd_parameters.to_csv(os.path.join(model_path,'model_parameters.txt'),index=False)

Expand Down Expand Up @@ -1815,7 +1833,7 @@ def train_animation_analyzer_onfly(self,data_path,model_path,out_path=None,dim=3
print('No train / validation folder!')


def train_combnet_onfly(self,data_path,model_path,out_path=None,dim_tconv=32,dim_conv=64,channel=1,time_step=15,level_tconv=1,level_conv=2,include_bodyparts=True,std=0,background_free=True,black_background=True,behavior_mode=0,social_distance=0):
def train_combnet_onfly(self,data_path,model_path,out_path=None,dim_tconv=32,dim_conv=64,channel=1,time_step=15,level_tconv=1,level_conv=2,include_bodyparts=True,std=0,background_free=True,black_background=True,behavior_mode=0,social_distance=0,color_costar=False):

# data_path: the folder that stores all the prepared training examples
# model_path: the path to the trained Animation Analyzer
Expand All @@ -1832,6 +1850,7 @@ def train_combnet_onfly(self,data_path,model_path,out_path=None,dim_tconv=32,dim
# black_background: whether to set background
# behavior_mode: 0--non-interactive, 1--interactive basic, 2--interactive advanced, 3--static images
# social_distance: a threshold (folds of size of a single animal) on whether to include individuals that are not main character in behavior examples
# color_costar: in 'interactive advanced' mode, whether to make the supporting roles RGB scale in animations

print('Training Categorizer with both Animation Analyzer and Pattern Recognizer using the behavior examples in: '+str(data_path))
self.log.append('Training Categorizer with both Animation Analyzer and Pattern Recognizer using the behavior examples in: '+str(data_path))
Expand Down Expand Up @@ -1870,7 +1889,12 @@ def train_combnet_onfly(self,data_path,model_path,out_path=None,dim_tconv=32,dim
else:
black_code=1

parameters={'classnames':list(train_data.classmapping.keys()),'dim_tconv':int(dim_tconv),'dim_conv':int(dim_conv),'channel':int(channel),'time_step':int(time_step),'network':2,'level_tconv':int(level_tconv),'level_conv':int(level_conv),'inner_code':int(inner_code),'std':int(std),'background_free':int(background_code),'black_background':int(black_code),'behavior_kind':int(behavior_mode),'social_distance':int(social_distance)}
if color_costar:
color_code=0
else:
color_code=1

parameters={'classnames':list(train_data.classmapping.keys()),'dim_tconv':int(dim_tconv),'dim_conv':int(dim_conv),'channel':int(channel),'time_step':int(time_step),'network':2,'level_tconv':int(level_tconv),'level_conv':int(level_conv),'inner_code':int(inner_code),'std':int(std),'background_free':int(background_code),'black_background':int(black_code),'behavior_kind':int(behavior_mode),'social_distance':int(social_distance),'color_code':int(color_code)}
pd_parameters=pd.DataFrame.from_dict(parameters)
pd_parameters.to_csv(os.path.join(model_path,'model_parameters.txt'),index=False)

Expand Down Expand Up @@ -1950,6 +1974,14 @@ def test_categorizer(self,groundtruth_path,model_path,result_path=None):
print('The behavior mode of the Categorizer: Interactive basic.')
elif behavior_mode==2:
print('The behavior mode of the Categorizer: Interactive advanced (Social distance '+str(parameters['social_distance'][0])+').')
if 'color_code' in list(parameters.keys()):
color_code=int(parameters['color_code'][0])
if color_code==0:
print('The Categorizer recognizes RGB scale main character and RGB scale supporting characters.')
else:
print('The Categorizer recognizes RGB scale main character and grayscale supporting characters.')
else:
print('The Categorizer recognizes RGB scale main character and grayscale supporting characters.')
else:
print('The behavior mode of the Categorizer: Static images (non-interactive).')
network=int(parameters['network'][0])
Expand Down
17 changes: 14 additions & 3 deletions LabGym/gui_analyzer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,7 @@ def __init__(self, parent):
self.black_background=True # whether to set background black
self.normalize_distance=True # whether to normalize the distance (in pixel) to the animal contour area
self.social_distance=0 # a threshold (folds of size of a single animal) on whether to include individuals that are not main character in behavior examples
self.color_costar=False # in 'interactive advanced' mode, whether to make the supporting roles RGB scale in animations
self.specific_behaviors={} # sex or identity-specific behaviors
self.correct_ID=False # whether to use sex or identity-specific behaviors to guide ID correction when ID switching is likely to happen

Expand Down Expand Up @@ -400,6 +401,12 @@ def select_categorizer(self,event):
self.social_distance=int(parameters['social_distance'][0])
if self.social_distance==0:
self.social_distance=float('inf')
if 'color_code' in parameters:
color_code=int(parameters['color_code'][0])
if color_code==0:
self.color_costar=True
else:
self.color_costar=False
self.text_detection.SetLabel('Only Detector-based detection method is available for the selected Categorizer.')
if self.behavior_mode==3:
self.text_detection.SetLabel('Only Detector-based detection method is available for the selected Categorizer.')
Expand Down Expand Up @@ -1034,7 +1041,7 @@ def analyze_behaviors(self,event):
if self.behavior_mode==1:
AAD.acquire_information_interact_basic(batch_size=self.detector_batch,background_free=self.background_free,black_background=self.black_background)
else:
AAD.acquire_information(batch_size=self.detector_batch,background_free=self.background_free,black_background=self.black_background)
AAD.acquire_information(batch_size=self.detector_batch,background_free=self.background_free,black_background=self.black_background,color_costar=self.color_costar)
if self.behavior_mode!=1:
AAD.craft_data()
if self.path_to_categorizer is not None:
Expand Down Expand Up @@ -1352,12 +1359,16 @@ def input_file(self,event):
dialog=wx.FileDialog(self,'Select the all_events.xlsx file.','',wildcard='all_events file (*.xlsx)|*.xlsx',style=wx.FD_OPEN)
if dialog.ShowModal()==wx.ID_OK:
all_events_file=Path(dialog.GetPath())
self.names_and_colors={}
names_and_colors={}
self.events_probability,self.time_points,behavior_names=parse_all_events_file(all_events_file)
colors=[('#ffffff',str(hex_code)) for hex_code in mpl.colors.cnames.values()]
for color,behavior in zip(colors,behavior_names):
self.names_and_colors[behavior]=color
names_and_colors[behavior]=color
self.text_inputfile.SetLabel(f'all_events.xlsx path: {all_events_file}')
if self.names_and_colors is None:
self.names_and_colors=names_and_colors
else:
self.names_and_colors=None
dialog.Destroy()


Expand Down
Loading