neuralcomputation/batik
0
1"""File for loading data into AnimalEditor"""
2import io
3from random import random
4from os.path import splitext
5from collections import OrderedDict
6import numpy as np
7from tempfile import NamedTemporaryFile
8
9from .annot import Annotations
10from .behavior import Behaviors
11
12def has_extension(fname:str, extension:str|list[str]) -> bool:
13 """
14 Checks to see if the passed in file name ends with an expected extension.
15 """
16 _, ext = splitext(fname)
17 if isinstance(extension, str):
18 return ext == extension
19 elif isinstance(extension, list):
20 return ext in extension
21
22def _clean_annotations(annotations):
23 """
24 While reading in behaviors from an .annot file, sometimes channels without normally
25 callable keys appear (i.e. keys that are strings which name a behavior), thus this
26 code only accepts keys which are strings.
27 """
28 if not annotations:
29 raise ValueError('No annotations found.')
30 clean_annot = OrderedDict()
31 for channel in annotations.keys():
32 channel_dict = OrderedDict()
33 for behavior_name in annotations[channel].keys():
34 if isinstance(behavior_name, str):
35 channel_dict.update({behavior_name : annotations[channel][behavior_name]})
36 clean_annot.update({channel: channel_dict})
37 return clean_annot
38
39def load_annot_sheet_txt(fname, offset = 0):
40 """
41 Generated a dictionary for retrieving the beginning and end frames of behaviors from
42 an .annot file.
43
44 Note that 0:00:00 is frame 1
45
46 Args:
47 fname - the path to the .annot file to be read (must be Caltech format)\n
48 offset - a value which offsets the start and end frame of each bout in
49 the sheet, as well as the absolute start and end frame of the file.
50 This value is optional, and is set to 0 by default
51
52 Returns:
53 annotations - dictionary of beginning and end frames for behaviors\n
54 start_time - the frame the movie started at (0:00:00 is 1)\n
55 end_time - the frame the movie ended at (0:00:00 is 1)\n
56 sample_rate - the sample rate reported within the file
57 """
58 # from bento for python
59 behaviors = Behaviors()
60 annot_sheet = Annotations(behaviors)
61 annot_sheet.read(fname)
62
63 sample_rate = annot_sheet.sample_rate()
64 annotations = OrderedDict()
65 for key in annot_sheet.channel_names():
66 annot_behaviors = OrderedDict()
67 bout_names = set()
68 for bout in annot_sheet.channel(key): #._bouts_by_start:
69 bout_names.add(bout.name())
70 for name in bout_names:
71 annot_behaviors.update({name : []})
72 for bout in annot_sheet.channel(key): #._bouts_by_start:
73 start_frame = bout.start().frames + offset
74 end_frame = bout.end().frames + offset
75 bout_frames = [start_frame, end_frame]
76 curr_table = annot_behaviors.get(bout.name())
77 new_table = curr_table.append(bout_frames)
78 annot_behaviors.update({bout.name : new_table})
79 for name in bout_names:
80 curr_table = annot_behaviors.get(name)
81 beh_array = np.array(curr_table)
82 annot_behaviors.update({name : beh_array})
83
84 annotations.update({key : annot_behaviors})
85 annotations = _clean_annotations(annotations)
86 start_time = annot_sheet.start_frame() + offset
87 end_time = annot_sheet.end_frame() + offset
88 return annotations, start_time, end_time, sample_rate
89
90def load_multiple_annotations(fnames):
91 """
92 Generates a single dictionary given multiple .annot files.
93 """
94 if not isinstance(fnames, list):
95 raise TypeError(f'Expected list[str], got {type(fnames)} instead.')
96 if not fnames:
97 raise ValueError('No file names passed in.')
98 if len(fnames) == 1:
99 return load_annot_sheet_txt(fnames[0])
100 head_annot, head_start_frame, head_end_frame, sample_rate = load_annot_sheet_txt(fnames[0])
101 end_frame = head_end_frame
102 for fname in fnames[1:]:
103 curr_annot, _, curr_end_frame, _ = load_annot_sheet_txt(fname, end_frame)
104 end_frame = curr_end_frame
105 for channel in curr_annot.keys():
106 if channel not in head_annot:
107 channel_dict = {}
108 head_annot.update({channel : channel_dict})
109 for behavior in curr_annot[channel].keys():
110 curr_behavior_bout_array = curr_annot[channel][behavior]
111 if channel in head_annot and behavior in head_annot[channel]:
112 new_bout_array = np.vstack((head_annot[channel][behavior],
113 curr_behavior_bout_array))
114 else:
115 new_bout_array = curr_behavior_bout_array
116 head_annot[channel].update({behavior : new_bout_array})
117 return head_annot, head_start_frame, end_frame, sample_rate
118
119def load_annot_sheet_txt_io(uploaded_file, offset = 0):
120 """
121 Generated a dictionary for retrieving the beginning and end frames of behaviors from
122 an .annot file.
123
124 Note that 0:00:00 is frame 1
125
126 Args:
127 fname - the path to the .annot file to be read (must be Caltech format)\n
128 offset - a value which offsets the start and end frame of each bout in
129 the sheet, as well as the absolute start and end frame of the file.
130 This value is optional, and is set to 0 by default
131
132 Returns:
133 annotations - dictionary of beginning and end frames for behaviors\n
134 start_time - the frame the movie started at (0:00:00 is 1)\n
135 end_time - the frame the movie ended at (0:00:00 is 1)\n
136 sample_rate - the sample rate reported within the file
137 """
138 # from bento for python
139 behaviors = Behaviors()
140 annot_sheet = Annotations(behaviors)
141
142 annot_sheet.read_io(uploaded_file)
143
144 sample_rate = annot_sheet.sample_rate()
145 annotations = OrderedDict()
146 for key in annot_sheet.channel_names():
147 annot_behaviors = OrderedDict()
148 bout_names = set()
149 for bout in annot_sheet.channel(key): #._bouts_by_start:
150 bout_names.add(bout.name())
151 for name in bout_names:
152 annot_behaviors.update({name : []})
153 for bout in annot_sheet.channel(key): #._bouts_by_start:
154 start_frame = bout.start().frames + offset
155 end_frame = bout.end().frames + offset
156 bout_frames = [start_frame, end_frame]
157 curr_table = annot_behaviors.get(bout.name())
158 new_table = curr_table.append(bout_frames)
159 annot_behaviors.update({bout.name : new_table})
160 for name in bout_names:
161 curr_table = annot_behaviors.get(name)
162 beh_array = np.array(curr_table)
163 annot_behaviors.update({name : beh_array})
164
165 annotations.update({key : annot_behaviors})
166 annotations = _clean_annotations(annotations)
167 start_time = annot_sheet.start_frame() + offset
168 end_time = annot_sheet.end_frame() + offset
169 return annotations, start_time, end_time, sample_rate
170
171def load_multiple_annotations_io(uploaded_files):
172 """
173 Generates a single dictionary given multiple .annot files.
174 """
175 if not isinstance(uploaded_files, list):
176 raise TypeError(f'Expected list, got {type(uploaded_files)} instead.')
177 if not uploaded_files:
178 raise ValueError('No file names passed in.')
179 if len(uploaded_files) == 1:
180 return load_annot_sheet_txt_io(uploaded_files[0])
181 head_annot, head_start_frame, head_end_frame, sample_rate = load_annot_sheet_txt_io(uploaded_files[0])
182 end_frame = head_end_frame
183 for uploaded_file in uploaded_files[1:]:
184 curr_annot, _, curr_end_frame, _ = load_annot_sheet_txt_io(uploaded_file, end_frame)
185 end_frame = curr_end_frame
186 for channel in curr_annot.keys():
187 if channel not in head_annot:
188 channel_dict = {}
189 head_annot.update({channel : channel_dict})
190 for behavior in curr_annot[channel].keys():
191 curr_behavior_bout_array = curr_annot[channel][behavior]
192 if channel in head_annot and behavior in head_annot[channel]:
193 new_bout_array = np.vstack((head_annot[channel][behavior],
194 curr_behavior_bout_array))
195 else:
196 new_bout_array = curr_behavior_bout_array
197 head_annot[channel].update({behavior : new_bout_array})
198 return head_annot, head_start_frame, end_frame, sample_rate
199