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 3 ctx1gfa

401 self.label = label 3 ctx1gfa

402 self.threshold = kwargs.get('threshold', 20) 3 ctx1gfa

403 self.upload = kwargs.get('upload', False) 3 ctx1gfa

404 self.twin = kwargs.get('twin', 150) 3 ctx1gfa

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

406 

407 self.load_data(sync=kwargs.get('sync', 'nidq'), location=kwargs.get('location', None)) 3 ctx1gfa

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): 3 ctx1gfa

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') 3 ctx1gfa

436 wheel_path = next(alf_path.rglob('*wheel.timestamps*')).parent 3 ctx1gfa

437 wheel = (fix_keys(alfio.load_object(wheel_path, 'wheel')) if location == 'SDSC' 2 ctx1fa

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

439 self.wheel_timestamps = wheel.timestamps 2 ctx1fa

440 # Compute interpolated wheel position and wheel times 

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

442 # Compute wheel velocity 

443 self.wheel_vel, _ = wh.velocity_filtered(wheel_pos, 1000) 2 ctx1fa

444 # Load in original camera times 

445 self.camera_path = str(next(self.session_path.joinpath('raw_video_data').glob(f'_iblrig_{self.label}Camera.raw*.mp4'))) 2 ctx1fa

446 self.camera_meta = vidio.get_video_meta(self.camera_path) 1 ctx1a

447 

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

449 if sync == 'nidq': 1 ctx1a

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

451 sync, chmap = get_sync_and_chn_map(self.session_path, sync_collection='raw_ephys_data') 1 ctx1a

452 sr = get_sync_fronts(sync, chmap[f'{self.label}_camera']) 1 ctx1a

453 self.ttls = sr.times[::2] 1 ctx1a

454 else: 

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

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

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

458 self.ttls = cam_extractor._times_from_bpod() 

459 

460 # Check if the ttl and video sizes match up 

461 self.tdiff = self.ttls.size - self.camera_meta['length'] 1 ctx1a

462 

463 # Load in original camera times if available otherwise set to ttls 

464 camera_times = next(alf_path.rglob(f'_ibl_{self.label}Camera.times*.npy'), None) 1 ctx1a

465 self.camera_times = alfio.load_file_content(camera_times) if camera_times else self.ttls 1 ctx1a

466 

467 if self.tdiff < 0: 1 ctx1a

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

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

470 self.ttl_times = self.ttls 

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

472 if self.camera_times.size != self.camera_meta['length']: 

473 self.camera_times = np.r_[self.camera_times, np.full((np.abs(self.tdiff)), np.nan)] 

474 self.short_flag = True 

475 elif self.tdiff > 0: 1 ctx1a

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

477 # tdiff ttls from the ttls 

478 self.ttl_times = self.ttls[self.tdiff:] 1 ctx1a

479 self.times = self.ttls[self.tdiff:] 1 ctx1a

480 if self.camera_times.size != self.camera_meta['length']: 1 ctx1a

481 self.camera_times = self.camera_times[self.tdiff:] 1 ctx1a

482 self.short_flag = False 1 ctx1a

483 

484 # Compute the frame rate of the camera 

485 self.frate = round(1 / np.nanmedian(np.diff(self.ttl_times))) 1 ctx1a

486 

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

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

489 try: 1 ctx1a

490 self.trials = alfio.load_file_content(next(alf_path.rglob('_ibl_trials.table*.pqt'))) 1 ctx1a

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

492 self.dlc = likelihood_threshold(self.dlc) 

493 self.behavior = True 

494 except (ALFObjectNotFound, StopIteration): 1 ctx1a

495 self.behavior = False 1 ctx1a

496 

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

498 self.frame_example = vidio.get_video_frames_preload(self.camera_path, np.arange(10, 11), mask=np.s_[:, :, 0]) 1 ctx1a

499 

500 def get_roi_mask(self): (empty)

501 """ 

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

503 use to compute the wheel motion energy 

504 :return: 

505 """ 

506 

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

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

509 else: 

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

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

512 

513 return roi, roi_mask 1 ctx1a

514 

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

516 """ 

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

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

519 background image 

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

521 :param thresold: threshold to differentiate artefact from background 

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

523 :return: mask of frames that are contaminated 

524 """ 

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

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

527 for idx, frame in enumerate(video_frames): 

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

529 high[idx] = ret 

530 

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

532 if normalise: 

533 high -= np.min(high) 

534 

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

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

537 

538 return contaminated_frames 

539 

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

541 """ 

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

543 in a parallel fashion jusing joblib.parallel 

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

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

546 :param wg: WindowGenerator 

547 :param iw: iteration of the WindowGenerator 

548 :return: 

549 """ 

550 

551 if iw == wg.nwin - 1: 

552 return 

553 

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

555 cap = cv2.VideoCapture(self.camera_path) 

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

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

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

559 

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

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

562 # this contamination interval 

563 if len(idx) != 0: 

564 

565 before_status = False 

566 after_status = False 

567 

568 counter = 0 

569 n_frames = 200 

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

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

572 # frame 

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

574 n_before_offset = (counter + 1) * n_frames 

575 first -= n_frames 

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

577 mask=self.mask) 

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

579 

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

581 before_status = True 

582 counter += 1 

583 if counter > 0: 

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

585 

586 counter = 0 

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

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

589 # frame 

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

591 n_after_offset = (counter + 1) * n_frames 

592 last += n_frames 

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

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

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

596 after_status = True 

597 counter += 1 

598 

599 if counter > 0: 

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

601 

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

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

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

605 for ints in intervals: 

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

607 ints = ints[1:] 

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

609 ints = ints[:-1] 

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

611 # We find all affected pixels 

612 for idx in ints: 

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

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

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

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

617 th_all += th 

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

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

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

621 for idx in ints: 

622 img = frames[idx] 

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

624 

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

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

627 if before_status: 

628 frames = frames[n_before_offset:]