diff --git a/ibllib/pipes/video_tasks.py b/ibllib/pipes/video_tasks.py index 3c1f4d19f..dd33b87b4 100644 --- a/ibllib/pipes/video_tasks.py +++ b/ibllib/pipes/video_tasks.py @@ -643,6 +643,7 @@ def _run(self, overwrite=True, run_qc=True, plot_qc=True): class LightningPose(base_tasks.VideoTask): # TODO: make one task per cam? + # TODO: separate pose and motion energy gpu = 1 io_charge = 100 level = 2 @@ -656,7 +657,9 @@ class LightningPose(base_tasks.VideoTask): def signature(self): signature = { 'input_files': [(f'_iblrig_{cam}Camera.raw.mp4', self.device_collection, True) for cam in self.cameras], - 'output_files': [(f'_ibl_{cam}Camera.lightningPose.pqt', 'alf', True) for cam in self.cameras] + 'output_files': [(f'_ibl_{cam}Camera.lightningPose.pqt', 'alf', True) for cam in self.cameras] + + [(f'{cam}Camera.ROIMotionEnergy.npy', 'alf', True) for cam in self.cameras] + + [(f'{cam}ROIMotionEnergy.position.npy', 'alf', True) for cam in self.cameras] } return signature @@ -723,6 +726,10 @@ def _run(self, overwrite=True, **kwargs): _logger.error(f"Corrupt raw video file {mp4_file}") self.status = -1 continue + + # --------------------------- + # Run pose estimation + # --------------------------- t0 = time.time() _logger.info(f'Running Lightning Pose on {label}Camera.') command2run = f"{self.scripts.joinpath('run_litpose.sh')} {str(self.env)} {mp4_file} {overwrite}" @@ -737,20 +744,61 @@ def _run(self, overwrite=True, **kwargs): info, error = process.communicate() if process.returncode != 0: error_str = error.decode("utf-8").strip() - _logger.error(f'Lightning pose failed for {label}Camera.\n\n' - f'++++++++ Output of subprocess for debugging ++++++++\n\n' - f'{error_str}\n' - f'++++++++++++++++++++++++++++++++++++++++++++\n') + _logger.error( + f'Lightning pose failed for {label}Camera.\n\n' + f'++++++++ Output of subprocess for debugging ++++++++\n\n' + f'{error_str}\n' + f'++++++++++++++++++++++++++++++++++++++++++++\n' + ) self.status = -1 + # We don't run motion energy, or add any files if LP failed to run continue else: _logger.info(f'{label} camera took {(time.time() - t0)} seconds') result = next(self.session_path.joinpath('alf').glob(f'_ibl_{label}Camera.lightningPose*.pqt')) actual_outputs.append(result) + # --------------------------- + # Run motion energy + # --------------------------- + t1 = time.time() + _logger.info(f'Computing motion energy for {label}Camera') + command2run = f"{self.scripts.joinpath('run_motion.sh')} {str(self.env)} {mp4_file} {result}" + _logger.info(command2run) + process = subprocess.Popen( + command2run, + shell=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + executable='/bin/bash', + ) + info, error = process.communicate() + if process.returncode != 0: + error_str = error.decode('utf-8').strip() + _logger.error( + f'Motion energy failed for {label}Camera.\n\n' + f'++++++++ Output of subprocess for debugging ++++++++\n\n' + f'{error_str}\n' + f'++++++++++++++++++++++++++++++++++++++++++++\n' + ) + self.status = -1 + continue + else: + _logger.info(f'{label} camera took {(time.time() - t1)} seconds') + actual_outputs.append(next(self.session_path.joinpath('alf').glob( + f'{label}Camera.ROIMotionEnergy*.npy'))) + actual_outputs.append(next(self.session_path.joinpath('alf').glob( + f'{label}ROIMotionEnergy.position*.npy'))) + except BaseException: _logger.error(traceback.format_exc()) self.status = -1 continue + # catch here if there are no raw videos present + if len(actual_outputs) == 0: + _logger.info('Did not find any videos for this session') + actual_outputs = None + self.status = -1 + return actual_outputs