Spaces:
Sleeping
Sleeping
| """File for loading data into AnimalEditor""" | |
| import io | |
| from random import random | |
| from os.path import splitext | |
| from collections import OrderedDict | |
| import numpy as np | |
| from tempfile import NamedTemporaryFile | |
| from .annot import Annotations | |
| from .behavior import Behaviors | |
| def has_extension(fname:str, extension:str|list[str]) -> bool: | |
| """ | |
| Checks to see if the passed in file name ends with an expected extension. | |
| """ | |
| _, ext = splitext(fname) | |
| if isinstance(extension, str): | |
| return ext == extension | |
| elif isinstance(extension, list): | |
| return ext in extension | |
| def _clean_annotations(annotations): | |
| """ | |
| While reading in behaviors from an .annot file, sometimes channels without normally | |
| callable keys appear (i.e. keys that are strings which name a behavior), thus this | |
| code only accepts keys which are strings. | |
| """ | |
| if not annotations: | |
| raise ValueError('No annotations found.') | |
| clean_annot = OrderedDict() | |
| for channel in annotations.keys(): | |
| channel_dict = OrderedDict() | |
| for behavior_name in annotations[channel].keys(): | |
| if isinstance(behavior_name, str): | |
| channel_dict.update({behavior_name : annotations[channel][behavior_name]}) | |
| clean_annot.update({channel: channel_dict}) | |
| return clean_annot | |
| def load_annot_sheet_txt(fname, offset = 0): | |
| """ | |
| Generated a dictionary for retrieving the beginning and end frames of behaviors from | |
| an .annot file. | |
| Note that 0:00:00 is frame 1 | |
| Args: | |
| fname - the path to the .annot file to be read (must be Caltech format)\n | |
| offset - a value which offsets the start and end frame of each bout in | |
| the sheet, as well as the absolute start and end frame of the file. | |
| This value is optional, and is set to 0 by default | |
| Returns: | |
| annotations - dictionary of beginning and end frames for behaviors\n | |
| start_time - the frame the movie started at (0:00:00 is 1)\n | |
| end_time - the frame the movie ended at (0:00:00 is 1)\n | |
| sample_rate - the sample rate reported within the file | |
| """ | |
| # from bento for python | |
| behaviors = Behaviors() | |
| annot_sheet = Annotations(behaviors) | |
| annot_sheet.read(fname) | |
| sample_rate = annot_sheet.sample_rate() | |
| annotations = OrderedDict() | |
| for key in annot_sheet.channel_names(): | |
| annot_behaviors = OrderedDict() | |
| bout_names = set() | |
| for bout in annot_sheet.channel(key): #._bouts_by_start: | |
| bout_names.add(bout.name()) | |
| for name in bout_names: | |
| annot_behaviors.update({name : []}) | |
| for bout in annot_sheet.channel(key): #._bouts_by_start: | |
| start_frame = bout.start().frames + offset | |
| end_frame = bout.end().frames + offset | |
| bout_frames = [start_frame, end_frame] | |
| curr_table = annot_behaviors.get(bout.name()) | |
| new_table = curr_table.append(bout_frames) | |
| annot_behaviors.update({bout.name : new_table}) | |
| for name in bout_names: | |
| curr_table = annot_behaviors.get(name) | |
| beh_array = np.array(curr_table) | |
| annot_behaviors.update({name : beh_array}) | |
| annotations.update({key : annot_behaviors}) | |
| annotations = _clean_annotations(annotations) | |
| start_time = annot_sheet.start_frame() + offset | |
| end_time = annot_sheet.end_frame() + offset | |
| return annotations, start_time, end_time, sample_rate | |
| def load_multiple_annotations(fnames): | |
| """ | |
| Generates a single dictionary given multiple .annot files. | |
| """ | |
| if not isinstance(fnames, list): | |
| raise TypeError(f'Expected list[str], got {type(fnames)} instead.') | |
| if not fnames: | |
| raise ValueError('No file names passed in.') | |
| if len(fnames) == 1: | |
| return load_annot_sheet_txt(fnames[0]) | |
| head_annot, head_start_frame, head_end_frame, sample_rate = load_annot_sheet_txt(fnames[0]) | |
| end_frame = head_end_frame | |
| for fname in fnames[1:]: | |
| curr_annot, _, curr_end_frame, _ = load_annot_sheet_txt(fname, end_frame) | |
| end_frame = curr_end_frame | |
| for channel in curr_annot.keys(): | |
| if channel not in head_annot: | |
| channel_dict = {} | |
| head_annot.update({channel : channel_dict}) | |
| for behavior in curr_annot[channel].keys(): | |
| curr_behavior_bout_array = curr_annot[channel][behavior] | |
| if channel in head_annot and behavior in head_annot[channel]: | |
| new_bout_array = np.vstack((head_annot[channel][behavior], | |
| curr_behavior_bout_array)) | |
| else: | |
| new_bout_array = curr_behavior_bout_array | |
| head_annot[channel].update({behavior : new_bout_array}) | |
| return head_annot, head_start_frame, end_frame, sample_rate | |
| def load_annot_sheet_txt_io(uploaded_file, offset = 0): | |
| """ | |
| Generated a dictionary for retrieving the beginning and end frames of behaviors from | |
| an .annot file. | |
| Note that 0:00:00 is frame 1 | |
| Args: | |
| fname - the path to the .annot file to be read (must be Caltech format)\n | |
| offset - a value which offsets the start and end frame of each bout in | |
| the sheet, as well as the absolute start and end frame of the file. | |
| This value is optional, and is set to 0 by default | |
| Returns: | |
| annotations - dictionary of beginning and end frames for behaviors\n | |
| start_time - the frame the movie started at (0:00:00 is 1)\n | |
| end_time - the frame the movie ended at (0:00:00 is 1)\n | |
| sample_rate - the sample rate reported within the file | |
| """ | |
| # from bento for python | |
| behaviors = Behaviors() | |
| annot_sheet = Annotations(behaviors) | |
| annot_sheet.read_io(uploaded_file) | |
| sample_rate = annot_sheet.sample_rate() | |
| annotations = OrderedDict() | |
| for key in annot_sheet.channel_names(): | |
| annot_behaviors = OrderedDict() | |
| bout_names = set() | |
| for bout in annot_sheet.channel(key): #._bouts_by_start: | |
| bout_names.add(bout.name()) | |
| for name in bout_names: | |
| annot_behaviors.update({name : []}) | |
| for bout in annot_sheet.channel(key): #._bouts_by_start: | |
| start_frame = bout.start().frames + offset | |
| end_frame = bout.end().frames + offset | |
| bout_frames = [start_frame, end_frame] | |
| curr_table = annot_behaviors.get(bout.name()) | |
| new_table = curr_table.append(bout_frames) | |
| annot_behaviors.update({bout.name : new_table}) | |
| for name in bout_names: | |
| curr_table = annot_behaviors.get(name) | |
| beh_array = np.array(curr_table) | |
| annot_behaviors.update({name : beh_array}) | |
| annotations.update({key : annot_behaviors}) | |
| annotations = _clean_annotations(annotations) | |
| start_time = annot_sheet.start_frame() + offset | |
| end_time = annot_sheet.end_frame() + offset | |
| return annotations, start_time, end_time, sample_rate | |
| def load_multiple_annotations_io(uploaded_files): | |
| """ | |
| Generates a single dictionary given multiple .annot files. | |
| """ | |
| if not isinstance(uploaded_files, list): | |
| raise TypeError(f'Expected list, got {type(uploaded_files)} instead.') | |
| if not uploaded_files: | |
| raise ValueError('No file names passed in.') | |
| if len(uploaded_files) == 1: | |
| return load_annot_sheet_txt_io(uploaded_files[0]) | |
| head_annot, head_start_frame, head_end_frame, sample_rate = load_annot_sheet_txt_io(uploaded_files[0]) | |
| end_frame = head_end_frame | |
| for uploaded_file in uploaded_files[1:]: | |
| curr_annot, _, curr_end_frame, _ = load_annot_sheet_txt_io(uploaded_file, end_frame) | |
| end_frame = curr_end_frame | |
| for channel in curr_annot.keys(): | |
| if channel not in head_annot: | |
| channel_dict = {} | |
| head_annot.update({channel : channel_dict}) | |
| for behavior in curr_annot[channel].keys(): | |
| curr_behavior_bout_array = curr_annot[channel][behavior] | |
| if channel in head_annot and behavior in head_annot[channel]: | |
| new_bout_array = np.vstack((head_annot[channel][behavior], | |
| curr_behavior_bout_array)) | |
| else: | |
| new_bout_array = curr_behavior_bout_array | |
| head_annot[channel].update({behavior : new_bout_array}) | |
| return head_annot, head_start_frame, end_frame, sample_rate | |