Merge pull request #28 from sidhulyalkar/new_facemap_integration_test

Update to create output dir during `FacemapPoseEstimation` processing
datajoint · Jan 5, 2024 · e8cbd7c · e8cbd7c
2 parents 2be11ec + 2089f26
commit e8cbd7c
Showing 1 changed file with 12 additions and 7 deletions.
diff --git a/element_facemap/facemap_inference.py b/element_facemap/facemap_inference.py
@@ -336,6 +336,15 @@ def make(self, key):
             "task_mode", "pose_estimation_output_dir"
         )
 
+        if not output_dir:
+            output_dir = FacemapPoseEstimationTask.infer_output_dir(
+                key, relative=True, mkdir=True
+            )
+            # update pose_estimation_output_dir
+            FacemapPoseEstimationTask.update1(
+                {**key, "pose_estimation_output_dir": output_dir.as_posix()}
+            )
+
         output_dir = find_full_path(fbe.get_facemap_root_data_dir(), output_dir)
         video_files = (FacemapPoseEstimationTask * fbe.VideoRecording.File & key).fetch(
             "file_path"
@@ -384,13 +393,9 @@ def make(self, key):
 
             bbox = (FacemapPoseEstimationTask & key).fetch1("bbox") or []
 
-            # Model Name of interest should be specified by user during facemap task params manual update
-            model_id = (FacemapPoseEstimationTask & key).fetch("model_id")
-
-            # Fetch model(.pt) file attachment to present working directory
-            facemap_model_name = (FacemapModel.File & f'model_id="{model_id}"').fetch1(
-                "model_file"
-            )
+            facemap_model_name = (
+                FacemapModel.File & f'model_id="{key["model_id"]}"'
+            ).fetch1("model_file")
 
             facemap_model_path = Path.cwd() / facemap_model_name
             models_root_dir = model_loader.get_models_dir()