Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add motion energy computation to LightningPose task #900

Draft
wants to merge 4 commits into
base: master
Choose a base branch
from
Draft

Add motion energy computation to LightningPose task #900

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4d208b4
add motion energy to LP task
themattinthehatt Dec 18, 2024
178e0f9
bugfix
themattinthehatt Dec 18, 2024
0b144d6
remove unnecessary comments
themattinthehatt Dec 18, 2024
d138737
flake
themattinthehatt Dec 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 53 additions & 5 deletions ibllib/pipes/video_tasks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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}"
Expand All @@ -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
Loading