Coverage for ibllib/io/extractors/video_motion.py: 45%

1""" 

2A module for aligning the wheel motion with the rotary encoder. Currently used by the camera QC 

3in order to check timestamp alignment. 

4""" 

5import matplotlib (empty)

6import matplotlib.pyplot as plt (empty)

7import matplotlib.gridspec as gridspec (empty)

8from matplotlib.widgets import RectangleSelector (empty)

9import numpy as np (empty)

10from scipy import signal, ndimage, interpolate (empty)

11import cv2 (empty)

12from itertools import cycle (empty)

13import matplotlib.animation as animation (empty)

14import logging (empty)

15from pathlib import Path (empty)

16from joblib import Parallel, delayed, cpu_count (empty)

17 

18from ibldsp.utils import WindowGenerator (empty)

19from one.api import ONE (empty)

20import ibllib.io.video as vidio (empty)

21from iblutil.util import Bunch (empty)

22from ibllib.io.extractors.ephys_fpga import get_sync_fronts, get_sync_and_chn_map (empty)

23import ibllib.io.raw_data_loaders as raw (empty)

24import ibllib.io.extractors.camera as cam (empty)

25from ibllib.plots.snapshot import ReportSnapshot (empty)

26import brainbox.video as video (empty)

27import brainbox.behavior.wheel as wh (empty)

28from brainbox.singlecell import bin_spikes (empty)

29from brainbox.behavior.dlc import likelihood_threshold, get_speed (empty)

30from brainbox.task.trials import find_trial_ids (empty)

31import one.alf.io as alfio (empty)

32from one.alf.exceptions import ALFObjectNotFound (empty)

33from one.alf.spec import is_session_path, is_uuid_string (empty)

34 

35 

36def find_nearest(array, value): (empty)

37 array = np.asarray(array) 3 ctx1deb

38 idx = (np.abs(array - value)).argmin() 3 ctx1deb

39 return idx 3 ctx1deb

40 

41 

42class MotionAlignment: (empty)

43 roi = {'left': ((800, 1020), (233, 1096)), 'right': ((426, 510), (104, 545)), 'body': ((402, 481), (31, 103))} (empty)

44 

45 def __init__(self, eid=None, one=None, log=logging.getLogger(__name__), stream=False, **kwargs): (empty)

46 self.one = one or ONE() 3 ctx1deb

47 self.eid = eid 3 ctx1deb

48 self.session_path = kwargs.pop('session_path', None) or self.one.eid2path(eid) 3 ctx1deb

49 self.ref = self.one.dict2ref(self.one.path2ref(self.session_path)) 3 ctx1deb

50 self.log = log 3 ctx1deb

51 self.trials = self.wheel = self.camera_times = None 3 ctx1deb

52 raw_cam_path = self.session_path.joinpath('raw_video_data') 3 ctx1deb

53 camera_path = list(raw_cam_path.glob('_iblrig_*Camera.raw.*')) 3 ctx1deb

54 if stream: 3 ctx1deb

55 self.video_paths = vidio.url_from_eid(self.eid) 

56 else: 

57 self.video_paths = {vidio.label_from_path(x): x for x in camera_path} 3 ctx1deb

58 self.data = Bunch() 3 ctx1deb

59 self.alignment = Bunch() 3 ctx1deb

60 

61 def align_all_trials(self, side='all'): (empty)

62 """Align all wheel motion for all trials""" 

63 if self.trials is None: 

64 self.load_data() 

65 if side == 'all': 

66 side = self.video_paths.keys() 

67 if not isinstance(side, str): 

68 # Try to iterate over sides 

69 [self.align_all_trials(s) for s in side] 

70 if side not in self.video_paths: 

71 raise ValueError(f'{side} camera video file not found') 

72 # Align each trial sequentially 

73 for i in np.arange(self.trials['intervals'].shape[0]): 

74 self.align_motion(i, display=False) 

75 

76 @staticmethod (empty)

77 def set_roi(video_path): (empty)

78 """Manually set the ROIs for a given set of videos 

79 TODO Improve docstring 

80 TODO A method for setting ROIs by label 

81 """ 

82 frame = vidio.get_video_frame(str(video_path), 0) 

83 

84 def line_select_callback(eclick, erelease): 

85 """ 

86 Callback for line selection. 

87 

88 *eclick* and *erelease* are the press and release events. 

89 """ 

90 x1, y1 = eclick.xdata, eclick.ydata 

91 x2, y2 = erelease.xdata, erelease.ydata 

92 print("(%3.2f, %3.2f) --> (%3.2f, %3.2f)" % (x1, y1, x2, y2)) 

93 return np.array([[x1, x2], [y1, y2]]) 

94 

95 plt.imshow(frame) 

96 roi = RectangleSelector(plt.gca(), line_select_callback, drawtype='box', useblit=True, button=[1, 3], 

97 # don't use middle button 

98 minspanx=5, minspany=5, spancoords='pixels', interactive=True) 

99 plt.show() 

100 ((x1, x2, *_), (y1, *_, y2)) = roi.corners 

101 col = np.arange(round(x1), round(x2), dtype=int) 

102 row = np.arange(round(y1), round(y2), dtype=int) 

103 return col, row 

104 

105 def load_data(self, download=False): (empty)

106 """ 

107 Load wheel, trial and camera timestamp data 

108 :return: wheel, trials 

109 """ 

110 if download: 

111 self.data.wheel = self.one.load_object(self.eid, 'wheel') 

112 self.data.trials = self.one.load_object(self.eid, 'trials') 

113 cam, det = self.one.load_datasets(self.eid, ['*Camera.times*']) 

114 self.data.camera_times = {vidio.label_from_path(d['rel_path']): ts for ts, d in zip(cam, det)} 

115 else: 

116 alf_path = self.session_path / 'alf' 

117 wheel_path = next(alf_path.rglob('*wheel.timestamps*')).parent 

118 self.data.wheel = alfio.load_object(wheel_path, 'wheel', short_keys=True) 

119 trials_path = next(alf_path.rglob('*trials.table*')).parent 

120 self.data.trials = alfio.load_object(trials_path, 'trials') 

121 self.data.camera_times = {vidio.label_from_path(x): alfio.load_file_content(x) for x in 

122 alf_path.rglob('*Camera.times*')} 

123 assert all(x is not None for x in self.data.values()) 

124 

125 def _set_eid_or_path(self, session_path_or_eid): (empty)

126 """Parse a given eID or session path 

127 If a session UUID is given, resolves and stores the local path and vice versa 

128 :param session_path_or_eid: A session eid or path 

129 :return: 

130 """ 

131 self.eid = None 

132 if is_uuid_string(str(session_path_or_eid)): 

133 self.eid = session_path_or_eid 

134 # Try to set session_path if data is found locally 

135 self.session_path = self.one.eid2path(self.eid) 

136 elif is_session_path(session_path_or_eid): 

137 self.session_path = Path(session_path_or_eid) 

138 if self.one is not None: 

139 self.eid = self.one.path2eid(self.session_path) 

140 if not self.eid: 

141 self.log.warning('Failed to determine eID from session path') 

142 else: 

143 self.log.error('Cannot run alignment: an experiment uuid or session path is required') 

144 raise ValueError("'session' must be a valid session path or uuid") 

145 

146 def align_motion(self, period=(-np.inf, np.inf), side='left', sd_thresh=10, display=False): (empty)

147 """ 

148 Align video to the wheel using cross-correlation of the video motion signal and the rotary 

149 encoder. 

150 

151 Parameters 

152 ---------- 

153 period : (float, float) 

154 The time period over which to do the alignment. 

155 side : {'left', 'right'} 

156 With which camera to perform the alignment. 

157 sd_thresh : float 

158 For plotting where the motion energy goes above this standard deviation threshold. 

159 display : bool 

160 When true, displays the aligned wheel motion energy along with the rotary encoder 

161 signal. 

162 

163 Returns 

164 ------- 

165 int 

166 Frame offset, i.e. by how many frames the video was shifted to match the rotary encoder 

167 signal. Negative values mean the video was shifted backwards with respect to the wheel 

168 timestamps. 

169 float 

170 The peak cross-correlation. 

171 numpy.ndarray 

172 The motion energy used in the cross-correlation, i.e. the frame difference for the 

173 period given. 

174 """ 

175 # Get data samples within period 

176 wheel = self.data['wheel'] 3 ctx1deb

177 self.alignment.label = side 3 ctx1deb

178 self.alignment.to_mask = lambda ts: np.logical_and(ts >= period[0], ts <= period[1]) 3 ctx1deb

179 camera_times = self.data['camera_times'][side] 3 ctx1deb

180 cam_mask = self.alignment.to_mask(camera_times) 3 ctx1deb

181 frame_numbers, = np.where(cam_mask) 3 ctx1deb

182 

183 if frame_numbers.size == 0: 3 ctx1deb

184 raise ValueError('No frames during given period') 1 ctx1b

185 

186 # Motion Energy 

187 camera_path = self.video_paths[side] 3 ctx1deb

188 roi = (*[slice(*r) for r in self.roi[side]], 0) 3 ctx1deb

189 try: 3 ctx1deb

190 # TODO Add function arg to make grayscale 

191 self.alignment.frames = vidio.get_video_frames_preload(camera_path, frame_numbers, mask=roi) 3 ctx1deb

192 assert self.alignment.frames.size != 0 3 ctx1deb

193 except AssertionError: 

194 self.log.error('Failed to open video') 

195 return None, None, None 

196 self.alignment.df, stDev = video.motion_energy(self.alignment.frames, 2) 3 ctx1deb

197 self.alignment.period = period # For plotting 3 ctx1deb

198 

199 # Calculate rotary encoder velocity trace 

200 x = camera_times[cam_mask] 3 ctx1deb

201 Fs = 1000 3 ctx1deb

202 pos, t = wh.interpolate_position(wheel.timestamps, wheel.position, freq=Fs) 3 ctx1deb

203 v, _ = wh.velocity_filtered(pos, Fs) 3 ctx1deb

204 interp_mask = self.alignment.to_mask(t) 3 ctx1deb

205 # Convert to normalized speed 

206 xs = np.unique([find_nearest(t[interp_mask], ts) for ts in x]) 3 ctx1deb

207 vs = np.abs(v[interp_mask][xs]) 3 ctx1deb

208 vs = (vs - np.min(vs)) / (np.max(vs) - np.min(vs)) 3 ctx1deb

209 

210 # FIXME This can be used as a goodness of fit measure 

211 USE_CV2 = False 3 ctx1deb

212 if USE_CV2: 3 ctx1deb

213 # convert from numpy format to openCV format 

214 dfCV = np.float32(self.alignment.df.reshape((-1, 1))) 

215 reCV = np.float32(vs.reshape((-1, 1))) 

216 

217 # perform cross correlation 

218 resultCv = cv2.matchTemplate(dfCV, reCV, cv2.TM_CCORR_NORMED) 

219 

220 # convert result back to numpy array 

221 xcorr = np.asarray(resultCv) 

222 else: 

223 xcorr = signal.correlate(self.alignment.df, vs) 3 ctx1deb

224 

225 # Cross correlate wheel speed trace with the motion energy 

226 CORRECTION = 2 3 ctx1deb

227 self.alignment.c = max(xcorr) 3 ctx1deb

228 self.alignment.xcorr = np.argmax(xcorr) 3 ctx1deb

229 self.alignment.dt_i = self.alignment.xcorr - xs.size + CORRECTION 3 ctx1deb

230 self.log.info(f'{side} camera, adjusted by {self.alignment.dt_i} frames') 3 ctx1deb

231 

232 if display: 3 ctx1deb

233 # Plot the motion energy 

234 fig, ax = plt.subplots(2, 1, sharex='all') 

235 y = np.pad(self.alignment.df, 1, 'edge') 

236 ax[0].plot(x, y, '-x', label='wheel motion energy') 

237 thresh = stDev > sd_thresh 

238 ax[0].vlines(x[np.array(np.pad(thresh, 1, 'constant', constant_values=False))], 0, 1, linewidth=0.5, linestyle=':', 

239 label=f'>{sd_thresh} s.d. diff') 

240 ax[1].plot(t[interp_mask], np.abs(v[interp_mask])) 

241 

242 # Plot other stuff 

243 dt = np.diff(camera_times[[0, np.abs(self.alignment.dt_i)]]) 

244 fps = 1 / np.diff(camera_times).mean() 

245 ax[0].plot(t[interp_mask][xs] - dt, vs, 'r-x', label='velocity (shifted)') 

246 ax[0].set_title('normalized motion energy, %s camera, %.0f fps' % (side, fps)) 

247 ax[0].set_ylabel('rate of change (a.u.)') 

248 ax[0].legend() 

249 ax[1].set_ylabel('wheel speed (rad / s)') 

250 ax[1].set_xlabel('Time (s)') 

251 

252 title = f'{self.ref}, from {period[0]:.1f}s - {period[1]:.1f}s' 

253 fig.suptitle(title, fontsize=16) 

254 fig.set_size_inches(19.2, 9.89) 

255 

256 return self.alignment.dt_i, self.alignment.c, self.alignment.df 3 ctx1deb

257 

258 def plot_alignment(self, energy=True, save=False): (empty)

259 if not self.alignment: 1 ctx1b

260 self.log.error('No alignment data, run `align_motion` first') 

261 return 

262 # Change backend based on save flag 

263 backend = matplotlib.get_backend().lower() 1 ctx1b

264 if (save and backend != 'agg') or (not save and backend == 'agg'): 1 ctx1b

265 new_backend = 'Agg' if save else 'Qt5Agg' 

266 self.log.warning('Switching backend from %s to %s', backend, new_backend) 

267 matplotlib.use(new_backend) 

268 from matplotlib import pyplot as plt 1 ctx1b

269 

270 # Main animated plots 

271 fig, axes = plt.subplots(nrows=2) 1 ctx1b

272 title = f'{self.ref}' # ', from {period[0]:.1f}s - {period[1]:.1f}s' 1 ctx1b

273 fig.suptitle(title, fontsize=16) 1 ctx1b

274 

275 wheel = self.data['wheel'] 1 ctx1b

276 wheel_mask = self.alignment['to_mask'](wheel.timestamps) 1 ctx1b

277 ts = self.data['camera_times'][self.alignment['label']] 1 ctx1b

278 frame_numbers, = np.where(self.alignment['to_mask'](ts)) 1 ctx1b

279 if energy: 1 ctx1b

280 self.alignment['frames'] = video.frame_diffs(self.alignment['frames'], 2) 1 ctx1b

281 frame_numbers = frame_numbers[1:-1] 1 ctx1b

282 data = {'frame_ids': frame_numbers} 1 ctx1b

283 

284 def init_plot(): 1 ctx1b

285 """ 

286 Plot the wheel data for the current trial 

287 :return: None 

288 """ 

289 data['im'] = axes[0].imshow(self.alignment['frames'][0]) 1 ctx1b

290 axes[0].axis('off') 1 ctx1b

291 axes[0].set_title(f'adjusted by {self.alignment["dt_i"]} frames') 1 ctx1b

292 

293 # Plot the wheel position 

294 ax = axes[1] 1 ctx1b

295 ax.clear() 1 ctx1b

296 ax.plot(wheel.timestamps[wheel_mask], wheel.position[wheel_mask], '-x') 1 ctx1b

297 

298 ts_0 = frame_numbers[0] 1 ctx1b

299 data['idx_0'] = ts_0 - self.alignment['dt_i'] 1 ctx1b

300 ts_0 = ts[ts_0 + self.alignment['dt_i']] 1 ctx1b

301 data['ln'] = ax.axvline(x=ts_0, color='k') 1 ctx1b

302 ax.set_xlim([ts_0 - (3 / 2), ts_0 + (3 / 2)]) 1 ctx1b

303 data['frame_num'] = 0 1 ctx1b

304 mkr = find_nearest(wheel.timestamps[wheel_mask], ts_0) 1 ctx1b

305 

306 data['marker'], = ax.plot(wheel.timestamps[wheel_mask][mkr], wheel.position[wheel_mask][mkr], 'r-x') 1 ctx1b

307 ax.set_ylabel('Wheel position (rad))') 1 ctx1b

308 ax.set_xlabel('Time (s))') 1 ctx1b

309 return 1 ctx1b

310 

311 def animate(i): 1 ctx1b

312 """ 

313 Callback for figure animation. Sets image data for current frame and moves pointer 

314 along axis 

315 :param i: unused; the current time step of the calling method 

316 :return: None 

317 """ 

318 if i < 0: 1 ctx1b

319 data['frame_num'] -= 1 

320 if data['frame_num'] < 0: 

321 data['frame_num'] = len(self.alignment['frames']) - 1 

322 else: 

323 data['frame_num'] += 1 1 ctx1b

324 if data['frame_num'] >= len(self.alignment['frames']): 1 ctx1b

325 data['frame_num'] = 0 1 ctx1b

326 i = data['frame_num'] # NB: This is index for current trial's frame list 1 ctx1b

327 

328 frame = self.alignment['frames'][i] 1 ctx1b

329 t_x = ts[data['idx_0'] + i] 1 ctx1b

330 data['ln'].set_xdata([t_x, t_x]) 1 ctx1b

331 axes[1].set_xlim([t_x - (3 / 2), t_x + (3 / 2)]) 1 ctx1b

332 data['im'].set_data(frame) 1 ctx1b

333 

334 mkr = find_nearest(wheel.timestamps[wheel_mask], t_x) 1 ctx1b

335 data['marker'].set_data([wheel.timestamps[wheel_mask][mkr]], [wheel.position[wheel_mask][mkr]]) 1 ctx1b

336 

337 return data['im'], data['ln'], data['marker'] 1 ctx1b

338 

339 anim = animation.FuncAnimation(fig, animate, init_func=init_plot, 1 ctx1b

340 frames=(range(len(self.alignment.df)) if save else cycle(range(60))), interval=20, 

341 blit=False, repeat=not save, cache_frame_data=False) 

342 anim.running = False 1 ctx1b

343 

344 def process_key(event): 1 ctx1b

345 """ 

346 Callback for key presses. 

347 :param event: a figure key_press_event 

348 :return: None 

349 """ 

350 if event.key.isspace(): 

351 if anim.running: 

352 anim.event_source.stop() 

353 else: 

354 anim.event_source.start() 

355 anim.running = ~anim.running 

356 elif event.key == 'right': 

357 if anim.running: 

358 anim.event_source.stop() 

359 anim.running = False 

360 animate(1) 

361 fig.canvas.draw() 

362 elif event.key == 'left': 

363 if anim.running: 

364 anim.event_source.stop() 

365 anim.running = False 

366 animate(-1) 

367 fig.canvas.draw() 

368 

369 fig.canvas.mpl_connect('key_press_event', process_key) 1 ctx1b

370 

371 # init_plot() 

372 # while True: 

373 # animate(0) 

374 if save: 1 ctx1b

375 filename = '%s_%c.mp4' % (self.ref, self.alignment['label'][0]) 1 ctx1b

376 if isinstance(save, (str, Path)): 1 ctx1b

377 filename = Path(save).joinpath(filename) 1 ctx1b

378 self.log.info(f'Saving to {filename}') 1 ctx1b

379 # Set up formatting for the movie files 

380 Writer = animation.writers['ffmpeg'] 1 ctx1b

381 writer = Writer(fps=24, metadata=dict(artist='Miles Wells'), bitrate=1800) 1 ctx1b

382 anim.save(str(filename), writer=writer) 1 ctx1b

383 else: 

384 plt.show() 

385 

386 

387class MotionAlignmentFullSession: (empty)

388 def __init__(self, session_path, label, **kwargs): (empty)

389 """ 

390 Class to extract camera times using video motion energy wheel alignment 

391 :param session_path: path of the session 

392 :param label: video label, only 'left' and 'right' videos are supported 

393 :param kwargs: threshold - the threshold to apply when identifying frames with artefacts (default 20) 

394 upload - whether to upload summary figure to alyx (default False) 

395 twin - the window length used when computing the shifts between the wheel and video 

396 nprocesses - the number of CPU processes to use 

397 sync - the type of sync scheme used (options 'nidq' or 'bpod') 

398 location - whether the code is being run on SDSC or not (options 'SDSC' or None) 

399 """ 

400 self.session_path = session_path 4 ctx1hfga

401 self.label = label 4 ctx1hfga

402 self.threshold = kwargs.get('threshold', 20) 4 ctx1hfga

403 self.upload = kwargs.get('upload', False) 4 ctx1hfga

404 self.twin = kwargs.get('twin', 150) 4 ctx1hfga

405 self.nprocess = kwargs.get('nprocess', int(cpu_count() - cpu_count() / 4)) 4 ctx1hfga

406 

407 self.load_data(sync=kwargs.get('sync', 'nidq'), location=kwargs.get('location', None)) 4 ctx1hfga

408 self.roi, self.mask = self.get_roi_mask() 1 ctx1a

409 

410 if self.upload: 1 ctx1a

411 self.one = ONE(mode='remote') 

412 self.one.alyx.authenticate() 

413 self.eid = self.one.path2eid(self.session_path) 

414 

415 def load_data(self, sync='nidq', location=None): (empty)

416 """ 

417 Loads relevant data from disk to perform motion alignment 

418 :param sync: type of sync used, 'nidq' or 'bpod' 

419 :param location: where the code is being run, if location='SDSC', the dataset uuids are removed 

420 when loading the data 

421 :return: 

422 """ 

423 def fix_keys(alf_object): 4 ctx1hfga

424 """ 

425 Given an alf object removes the dataset uuid from the keys 

426 :param alf_object: 

427 :return: 

428 """ 

429 ob = Bunch() 

430 for key in alf_object.keys(): 

431 vals = alf_object[key] 

432 ob[key.split('.')[0]] = vals 

433 return ob 

434 

435 alf_path = self.session_path.joinpath('alf') 4 ctx1hfga

436 wheel_path = next(alf_path.rglob('*wheel.timestamps*')).parent 4 ctx1hfga

437 wheel = (fix_keys(alfio.load_object(wheel_path, 'wheel')) if location == 'SDSC' 3 ctx1fga

438 else alfio.load_object(wheel_path, 'wheel')) 

439 self.wheel_timestamps = wheel.timestamps 3 ctx1fga

440 # Compute interpolated wheel position and wheel times 

441 wheel_pos, self.wheel_time = wh.interpolate_position(wheel.timestamps, wheel.position, freq=1000) 3 ctx1fga

442 # Compute wheel velocity 

443 self.wheel_vel, _ = wh.velocity_filtered(wheel_pos, 1000) 3 ctx1fga

444 # Load in original camera times 

445 self.camera_times = alfio.load_file_content(next(alf_path.rglob(f'_ibl_{self.label}Camera.times*.npy'))) 3 ctx1fga

446 self.camera_path = str(next(self.session_path.joinpath('raw_video_data').glob(f'_iblrig_{self.label}Camera.raw*.mp4'))) 3 ctx1fga

447 self.camera_meta = vidio.get_video_meta(self.camera_path) 2 ctx1fa

448 

449 # TODO should read in the description file to get the correct sync location 

450 if sync == 'nidq': 2 ctx1fa

451 # If the sync is 'nidq' we read in the camera ttls from the spikeglx sync object 

452 sync, chmap = get_sync_and_chn_map(self.session_path, sync_collection='raw_ephys_data') 2 ctx1fa

453 sr = get_sync_fronts(sync, chmap[f'{self.label}_camera']) 2 ctx1fa

454 self.ttls = sr.times[::2] 2 ctx1fa

455 else: 

456 # Otherwise we assume the sync is 'bpod' and we read in the camera ttls from the raw bpod data 

457 cam_extractor = cam.CameraTimestampsBpod(session_path=self.session_path) 

458 cam_extractor.bpod_trials = raw.load_data(self.session_path, task_collection='raw_behavior_data') 

459 self.ttls = cam_extractor._times_from_bpod() 

460 

461 # Check if the ttl and video sizes match up 

462 self.tdiff = self.ttls.size - self.camera_meta['length'] 2 ctx1fa

463 

464 if self.tdiff < 0: 2 ctx1fa

465 # In this case there are fewer ttls than camera frames. This is not ideal, for now we pad the ttls with 

466 # nans but if this is too many we reject the wheel alignment based on the qc 

467 self.ttl_times = self.ttls 

468 self.times = np.r_[self.ttl_times, np.full((np.abs(self.tdiff)), np.nan)] 

469 self.short_flag = True 

470 elif self.tdiff > 0: 2 ctx1fa

471 # In this case there are more ttls than camera frames. This happens often, for now we remove the first 

472 # tdiff ttls from the ttls 

473 self.ttl_times = self.ttls[self.tdiff:] 2 ctx1fa

474 self.times = self.ttls[self.tdiff:] 2 ctx1fa

475 self.short_flag = False 2 ctx1fa

476 

477 # Compute the frame rate of the camera 

478 self.frate = round(1 / np.nanmedian(np.diff(self.ttl_times))) 2 ctx1fa

479 

480 # We attempt to load in some behavior data (trials and dlc). This is only needed for the summary plots, having 

481 # trial aligned paw velocity (from the dlc) is a nice sanity check to make sure the alignment went well 

482 try: 2 ctx1fa

483 self.trials = alfio.load_file_content(next(alf_path.rglob('_ibl_trials.table*.pqt'))) 2 ctx1fa

484 self.dlc = alfio.load_file_content(next(alf_path.rglob(f'_ibl_{self.label}Camera.dlc*.pqt'))) 

485 self.dlc = likelihood_threshold(self.dlc) 

486 self.behavior = True 

487 except (ALFObjectNotFound, StopIteration): 2 ctx1fa

488 self.behavior = False 2 ctx1fa

489 

490 # Load in a single frame that we will use for the summary plot 

491 self.frame_example = vidio.get_video_frames_preload(self.camera_path, np.arange(10, 11), mask=np.s_[:, :, 0]) 2 ctx1fa

492 

493 def get_roi_mask(self): (empty)

494 """ 

495 Compute the region of interest mask for a given camera. This corresponds to a box in the video that we will 

496 use to compute the wheel motion energy 

497 :return: 

498 """ 

499 

500 if self.label == 'right': 1 ctx1a

501 roi = ((450, 512), (120, 200)) 1 ctx1a

502 else: 

503 roi = ((900, 1024), (850, 1010)) 

504 roi_mask = (*[slice(*r) for r in roi], 0) 1 ctx1a

505 

506 return roi, roi_mask 1 ctx1a

507 

508 def find_contaminated_frames(self, video_frames, thresold=20, normalise=True): (empty)

509 """ 

510 Finds frames in the video that have artefacts such as the mouse's paw or a human hand. In order to determine 

511 frames with contamination an Otsu thresholding is applied to each frame to detect the artefact from the 

512 background image 

513 :param video_frames: np array of video frames (nframes, nwidth, nheight) 

514 :param thresold: threshold to differentiate artefact from background 

515 :param normalise: whether to normalise the threshold values for each frame to the baseline 

516 :return: mask of frames that are contaminated 

517 """ 

518 high = np.zeros((video_frames.shape[0])) 

519 # Iterate through each frame and compute and store the otsu threshold value for each frame 

520 for idx, frame in enumerate(video_frames): 

521 ret, _ = cv2.threshold(cv2.GaussianBlur(frame, (5, 5), 0), 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) 

522 high[idx] = ret 

523 

524 # If normalise is True, we divide the threshold values for each frame by the minimum value 

525 if normalise: 

526 high -= np.min(high) 

527 

528 # Identify the frames that have a threshold value greater than the specified threshold cutoff 

529 contaminated_frames = np.where(high > thresold)[0] 

530 

531 return contaminated_frames 

532 

533 def compute_motion_energy(self, first, last, wg, iw): (empty)

534 """ 

535 Computes the video motion energy for frame indexes between first and last. This function is written to be run 

536 in a parallel fashion jusing joblib.parallel 

537 :param first: first frame index of frame interval to consider 

538 :param last: last frame index of frame interval to consider 

539 :param wg: WindowGenerator 

540 :param iw: iteration of the WindowGenerator 

541 :return: 

542 """ 

543 

544 if iw == wg.nwin - 1: 

545 return 

546 

547 # Open the video and read in the relvant video frames between first idx and last idx 

548 cap = cv2.VideoCapture(self.camera_path) 

549 frames = vidio.get_video_frames_preload(cap, np.arange(first, last), mask=self.mask) 

550 # Identify if any of the frames have artefacts in them 

551 idx = self.find_contaminated_frames(frames, self.threshold) 

552 

553 # If some of the frames are contaminated we find all the continuous intervals of contamination 

554 # and set the value for contaminated pixels for these frames to the average of the first frame before and after 

555 # this contamination interval 

556 if len(idx) != 0: 

557 

558 before_status = False 

559 after_status = False 

560 

561 counter = 0 

562 n_frames = 200 

563 # If it is the first frame that is contaminated, we need to read in a bit more of the video to find a 

564 # frame prior to contamination. We attempt this 20 times, after that we just take the value for the first 

565 # frame 

566 while np.any(idx == 0) and counter < 20 and iw != 0: 

567 n_before_offset = (counter + 1) * n_frames 

568 first -= n_frames 

569 extra_frames = vidio.get_video_frames_preload(cap, frame_numbers=np.arange(first - n_frames, first), 

570 mask=self.mask) 

571 frames = np.concatenate([extra_frames, frames], axis=0) 

572 

573 idx = self.find_contaminated_frames(frames, self.threshold) 

574 before_status = True 

575 counter += 1 

576 if counter > 0: 

577 print(f'In before: {counter}') 

578 

579 counter = 0 

580 # If it is the last frame that is contaminated, we need to read in a bit more of the video to find a 

581 # frame after the contamination. We attempt this 20 times, after that we just take the value for the last 

582 # frame 

583 while np.any(idx == frames.shape[0] - 1) and counter < 20 and iw != wg.nwin - 1: 

584 n_after_offset = (counter + 1) * n_frames 

585 last += n_frames 

586 extra_frames = vidio.get_video_frames_preload(cap, frame_numbers=np.arange(last, last + n_frames), mask=self.mask) 

587 frames = np.concatenate([frames, extra_frames], axis=0) 

588 idx = self.find_contaminated_frames(frames, self.threshold) 

589 after_status = True 

590 counter += 1 

591 

592 if counter > 0: 

593 print(f'In after: {counter}') 

594 

595 # We find all the continuous intervals that contain contamination and fix the affected pixels 

596 # by taking the average value of the frame prior and after contamination 

597 intervals = np.split(idx, np.where(np.diff(idx) != 1)[0] + 1) 

598 for ints in intervals: 

599 if len(ints) > 0 and ints[0] == 0: 

600 ints = ints[1:] 

601 if len(ints) > 0 and ints[-1] == frames.shape[0] - 1: 

602 ints = ints[:-1] 

603 th_all = np.zeros_like(frames[0]) 

604 # We find all affected pixels 

605 for idx in ints: 

606 img = np.copy(frames[idx]) 

607 blur = cv2.GaussianBlur(img, (5, 5), 0) 

608 ret, th = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) 

609 th = cv2.GaussianBlur(th, (5, 5), 10) 

610 th_all += th 

611 # Compute the average image of the frame prior and after the interval 

612 vals = np.mean(np.dstack([frames[ints[0] - 1], frames[ints[-1] + 1]]), axis=-1) 

613 # For each frame set the affected pixels to the value of the clean average image 

614 for idx in ints: 

615 img = frames[idx] 

616 img[th_all > 0] = vals[th_all > 0] 

617 

618 # If we have read in extra video frames we need to cut these off and make sure we only 

619 # consider the frames between the interval first and last given as args 

620 if before_status: 

621 frames = frames[n_before_offset:] 

622 if after_status: 

623 frames = frames[:(-1 * n_after_offset)] 

624 

625 # Once the frames have been cleaned we compute the motion energy between frames 

626 frame_me, _ = video.motion_energy(frames, diff=2, normalize=False) 

627 

628 cap.release() 

629 

630 return frame_me[2:] 

631 

632 def compute_shifts(self, times, me, first, last, iw, wg): (empty)