From eecb74cb5c94cf125fa9e9f57d34899ea2444179 Mon Sep 17 00:00:00 2001 From: Thinh Nguyen Date: Wed, 7 Aug 2024 16:57:34 -0500 Subject: [PATCH 1/6] feat: add `PoseEstimationReport` table --- element_deeplabcut/model.py | 82 ++++++++++++++++++++++++++++++++++--- 1 file changed, 76 insertions(+), 6 deletions(-) diff --git a/element_deeplabcut/model.py b/element_deeplabcut/model.py index 08b1a9a..adc5096 100644 --- a/element_deeplabcut/model.py +++ b/element_deeplabcut/model.py @@ -705,7 +705,7 @@ class BodyPartPosition(dj.Part): def make(self, key): """.populate() method will launch training for each PoseEstimationTask""" # ID model and directories - dlc_model = (Model & key).fetch1() + dlc_model_ = (Model & key).fetch1() task_mode, output_dir = (PoseEstimationTask & key).fetch1( "task_mode", "pose_estimation_output_dir" ) @@ -726,7 +726,7 @@ def make(self, key): # - video_filepaths: full paths to the video files for inference # - analyze_video_params: optional parameters to analyze video project_path = find_full_path( - get_dlc_root_data_dir(), dlc_model["project_path"] + get_dlc_root_data_dir(), dlc_model_["project_path"] ) video_relpaths = list((VideoRecording.File & key).fetch("file_path")) video_filepaths = [ @@ -740,9 +740,9 @@ def make(self, key): @memoized_result( uniqueness_dict={ **analyze_video_params, - "project_path": dlc_model["project_path"], - "shuffle": dlc_model["shuffle"], - "trainingsetindex": dlc_model["trainingsetindex"], + "project_path": dlc_model_["project_path"], + "shuffle": dlc_model_["shuffle"], + "trainingsetindex": dlc_model_["trainingsetindex"], "video_filepaths": video_relpaths, }, output_directory=output_dir, @@ -751,7 +751,7 @@ def do_pose_estimation(): dlc_reader.do_pose_estimation( key, video_filepaths, - dlc_model, + dlc_model_, project_path, output_dir, **analyze_video_params, @@ -819,6 +819,76 @@ def get_trajectory(cls, key: dict, body_parts: list = "all") -> pd.DataFrame: return df +@schema +class PoseEstimationReport(dj.Computed): + definition = """ + -> PoseEstimation + """ + + class LabeledVideo(dj.Part): + definition = """ + -> master + -> VideoRecording.File + --- + labeled_video_path: varchar(255) + """ + + @property + def key_source(self): + return PoseEstimation & RecordingInfo + + def make(self, key): + import deeplabcut + + # some default settings + outputframerate = 5 # final labeled video will be 5 Hz + + dlc_model_ = (Model & key).fetch1() + fps, nframes = (RecordingInfo & key).fetch1("fps", "nframes") + output_dir = (PoseEstimationTask & key).fetch1("pose_estimation_output_dir") + output_dir = find_full_path(get_dlc_root_data_dir(), output_dir) + + project_path = find_full_path( + get_dlc_root_data_dir(), dlc_model_["project_path"] + ) + dlc_config = project_path / "dj_dlc_config.yaml" + + entries = [] + for vkey in (VideoRecording.File & key).fetch("KEY"): + video_file = (VideoRecording.File & vkey).fetch1("file_path") + video_file = find_full_path(get_dlc_root_data_dir(), video_file) + + deeplabcut.create_labeled_video( + config=dlc_config.as_posix(), + videos=[video_file.as_posix()], + shuffle=dlc_model_["shuffle"], + trainingsetindex=dlc_model_["trainingsetindex"], + destfolder=output_dir, + Frames2plot=np.arange(0, nframes, int(fps / outputframerate)), + outputframerate=outputframerate, + displaycropped=True, + draw_skeleton=True, + save_frames=False, + ) + + labeled_video_path = next( + output_dir.glob(f"{video_file.stem}*_labeled.mp4") + ) + + entries.append( + { + **key, + **vkey, + "labeled_video_path": labeled_video_path.relative_to( + get_dlc_processed_data_dir() + ).as_posix(), + } + ) + + self.insert1(key) + self.LabeledVideo.insert(entries) + + def str_to_bool(value) -> bool: """Return whether the provided string represents true. Otherwise false. From 893399bccdde472155dcd5dd29072ccd695a77d5 Mon Sep 17 00:00:00 2001 From: Thinh Nguyen Date: Thu, 8 Aug 2024 11:10:04 -0500 Subject: [PATCH 2/6] fix: displaycropped=False --- element_deeplabcut/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/element_deeplabcut/model.py b/element_deeplabcut/model.py index f364660..23bafce 100644 --- a/element_deeplabcut/model.py +++ b/element_deeplabcut/model.py @@ -889,7 +889,7 @@ def make(self, key): destfolder=output_dir, Frames2plot=np.arange(0, nframes, int(fps / outputframerate)), outputframerate=outputframerate, - displaycropped=True, + displaycropped=False, draw_skeleton=True, save_frames=False, ) From 323b8fdbd64036a1d456e5c77c26fe5a9f2d9f49 Mon Sep 17 00:00:00 2001 From: Thinh Nguyen Date: Thu, 8 Aug 2024 13:23:16 -0500 Subject: [PATCH 3/6] feat: allow parameterization of the `PoseEstimationReport` step --- element_deeplabcut/model.py | 69 ++++++++++++++++++++++++++++++------- 1 file changed, 57 insertions(+), 12 deletions(-) diff --git a/element_deeplabcut/model.py b/element_deeplabcut/model.py index 481b8a4..d065e94 100644 --- a/element_deeplabcut/model.py +++ b/element_deeplabcut/model.py @@ -733,10 +733,16 @@ def make(self, key): find_full_path(get_dlc_root_data_dir(), fp).as_posix() for fp in video_relpaths ] - analyze_video_params = (PoseEstimationTask & key).fetch1( + pose_estimation_params = (PoseEstimationTask & key).fetch1( "pose_estimation_params" ) or {} + # expect a nested dictionary with "analyze_videos" params + # if not, assume "pose_estimation_params" as a flat dictionary that include relevant "analyze_videos" params + analyze_video_params = ( + pose_estimation_params.get("analyze_videos") or pose_estimation_params + ) + @memoized_result( uniqueness_dict={ **analyze_video_params, @@ -885,6 +891,20 @@ def key_source(self): def make(self, key): import deeplabcut + pose_estimation_params = (PoseEstimationTask & key).fetch1( + "pose_estimation_params" + ) or {} + + # expect a nested dictionary with "create_labeled_video" and "extract_outlier_frames" params + # if not, assume "pose_estimation_params" as a flat dictionary + create_labeled_video_params = ( + pose_estimation_params.get("create_labeled_video") or pose_estimation_params + ) + extract_outlier_frames_params = ( + pose_estimation_params.get("extract_outlier_frames") + or pose_estimation_params + ) + # some default settings outputframerate = 5 # final labeled video will be 5 Hz @@ -903,18 +923,43 @@ def make(self, key): video_file = (VideoRecording.File & vkey).fetch1("file_path") video_file = find_full_path(get_dlc_root_data_dir(), video_file) - deeplabcut.create_labeled_video( - config=dlc_config.as_posix(), - videos=[video_file.as_posix()], - shuffle=dlc_model_["shuffle"], - trainingsetindex=dlc_model_["trainingsetindex"], - destfolder=output_dir, - Frames2plot=np.arange(0, nframes, int(fps / outputframerate)), - outputframerate=outputframerate, - displaycropped=False, - draw_skeleton=True, - save_frames=False, + # -- create labeled video -- + create_labeled_video_kwargs = { + k: v + for k, v in create_labeled_video_params.items() + if k in inspect.signature(deeplabcut.create_labeled_video).parameters + } + create_labeled_video_kwargs.update( + dict( + config=dlc_config.as_posix(), + videos=[video_file.as_posix()], + shuffle=dlc_model_["shuffle"], + trainingsetindex=dlc_model_["trainingsetindex"], + modelprefix=dlc_model_["model_prefix"], + destfolder=output_dir, + Frames2plot=np.arange(0, nframes, int(fps / outputframerate)), + outputframerate=outputframerate, + ) + ) + deeplabcut.create_labeled_video(**create_labeled_video_kwargs) + + # -- extract outlier frames -- + extract_outlier_frames_kwargs = { + k: v + for k, v in extract_outlier_frames_params.items() + if k in inspect.signature(deeplabcut.extract_outlier_frames).parameters + } + extract_outlier_frames_kwargs.update( + dict( + config=dlc_config.as_posix(), + videos=[video_file.as_posix()], + shuffle=dlc_model_["shuffle"], + trainingsetindex=dlc_model_["trainingsetindex"], + modelprefix=dlc_model_["model_prefix"], + destfolder=output_dir, + ) ) + deeplabcut.extract_outlier_frames(**create_labeled_video_kwargs) labeled_video_path = next( output_dir.glob(f"{video_file.stem}*_labeled.mp4") From ea3c335c8a4c2593b00adbb58a118c6a6c300090 Mon Sep 17 00:00:00 2001 From: Thinh Nguyen Date: Thu, 8 Aug 2024 14:56:12 -0500 Subject: [PATCH 4/6] fix: run `create_labeled_video` only --- element_deeplabcut/model.py | 38 ++++++++----------------------------- 1 file changed, 8 insertions(+), 30 deletions(-) diff --git a/element_deeplabcut/model.py b/element_deeplabcut/model.py index d065e94..7eb742e 100644 --- a/element_deeplabcut/model.py +++ b/element_deeplabcut/model.py @@ -871,17 +871,17 @@ def get_trajectory(cls, key: dict, body_parts: list = "all") -> pd.DataFrame: @schema -class PoseEstimationReport(dj.Computed): +class LabeledVideo(dj.Computed): definition = """ -> PoseEstimation """ - class LabeledVideo(dj.Part): + class File(dj.Part): definition = """ -> master -> VideoRecording.File --- - labeled_video_path: varchar(255) + labeled_video_path: varchar(255) # relative path to labeled video """ @property @@ -900,13 +900,10 @@ def make(self, key): create_labeled_video_params = ( pose_estimation_params.get("create_labeled_video") or pose_estimation_params ) - extract_outlier_frames_params = ( - pose_estimation_params.get("extract_outlier_frames") - or pose_estimation_params - ) - # some default settings - outputframerate = 5 # final labeled video will be 5 Hz + outputframerate = create_labeled_video_params.pop( + "outputframerate", 5 + ) # final labeled video FPS defaults to 5 Hz dlc_model_ = (Model & key).fetch1() fps, nframes = (RecordingInfo & key).fetch1("fps", "nframes") @@ -936,35 +933,16 @@ def make(self, key): shuffle=dlc_model_["shuffle"], trainingsetindex=dlc_model_["trainingsetindex"], modelprefix=dlc_model_["model_prefix"], - destfolder=output_dir, + destfolder=output_dir.as_posix(), Frames2plot=np.arange(0, nframes, int(fps / outputframerate)), outputframerate=outputframerate, ) ) deeplabcut.create_labeled_video(**create_labeled_video_kwargs) - # -- extract outlier frames -- - extract_outlier_frames_kwargs = { - k: v - for k, v in extract_outlier_frames_params.items() - if k in inspect.signature(deeplabcut.extract_outlier_frames).parameters - } - extract_outlier_frames_kwargs.update( - dict( - config=dlc_config.as_posix(), - videos=[video_file.as_posix()], - shuffle=dlc_model_["shuffle"], - trainingsetindex=dlc_model_["trainingsetindex"], - modelprefix=dlc_model_["model_prefix"], - destfolder=output_dir, - ) - ) - deeplabcut.extract_outlier_frames(**create_labeled_video_kwargs) - labeled_video_path = next( output_dir.glob(f"{video_file.stem}*_labeled.mp4") ) - entries.append( { **key, @@ -976,7 +954,7 @@ def make(self, key): ) self.insert1(key) - self.LabeledVideo.insert(entries) + self.File.insert(entries) def str_to_bool(value) -> bool: From 4ddfb34579660c6c31853b40d1c858cb8e841132 Mon Sep 17 00:00:00 2001 From: Thinh Nguyen Date: Thu, 8 Aug 2024 15:07:46 -0500 Subject: [PATCH 5/6] fix: robust search for dj_dlc_config*.yaml --- element_deeplabcut/model.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/element_deeplabcut/model.py b/element_deeplabcut/model.py index fc4505f..61341e5 100644 --- a/element_deeplabcut/model.py +++ b/element_deeplabcut/model.py @@ -19,6 +19,8 @@ from .readers import dlc_reader schema = dj.schema() +logger = dj.logger + _linking_module = None @@ -916,7 +918,16 @@ def make(self, key): project_path = find_full_path( get_dlc_root_data_dir(), dlc_model_["project_path"] ) - dlc_config = project_path / "dj_dlc_config.yaml" + + try: + dlc_config = next(output_dir.glob("dj_dlc_config*.yaml")) + dlc_config = project_path / dlc_config.name + assert dlc_config.exists() + except (StopIteration, AssertionError): + dlc_config = next(project_path.glob("dj_dlc_config*.yaml")) + logger.warning( + f"No dj_dlc_config*.yaml file found in {output_dir} - this is unexpected.\nUsing {dlc_config}" + ) entries = [] for vkey in (VideoRecording.File & key).fetch("KEY"): From 8c5a5e9f445419207a477d2cb174ab3c9f62c205 Mon Sep 17 00:00:00 2001 From: Thinh Nguyen Date: Mon, 12 Aug 2024 16:02:48 -0500 Subject: [PATCH 6/6] feat(LabeledVideo): add `labeled_video_file` attribute as `filepath` --- element_deeplabcut/model.py | 2 ++ setup.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/element_deeplabcut/model.py b/element_deeplabcut/model.py index 8720b41..6609fd2 100644 --- a/element_deeplabcut/model.py +++ b/element_deeplabcut/model.py @@ -931,6 +931,7 @@ class File(dj.Part): -> VideoRecording.File --- labeled_video_path: varchar(255) # relative path to labeled video + labeled_video_file=null: filepath@dlc-processed """ @property @@ -1008,6 +1009,7 @@ def make(self, key): "labeled_video_path": labeled_video_path.relative_to( get_dlc_processed_data_dir() ).as_posix(), + "labeled_video_file": labeled_video_path.as_posix(), } ) diff --git a/setup.py b/setup.py index 368a49b..ba455c5 100644 --- a/setup.py +++ b/setup.py @@ -48,6 +48,6 @@ "tests": ["pytest", "pytest-cov", "shutils"], "dlc-pytorch": [ "deeplabcut @ git+https://github.com/DeepLabCut/DeepLabCut.git@pytorch_dlc" - ] + ], }, )