From 475ca271ef6f6b33087459bb446ba746e599c049 Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@datajoint.com>
Date: Fri, 23 Aug 2024 09:30:27 -0500
Subject: [PATCH] chore: code cleanup

---
 element_facemap/facemap_inference.py | 64 ++++++++++------------------
 1 file changed, 23 insertions(+), 41 deletions(-)

diff --git a/element_facemap/facemap_inference.py b/element_facemap/facemap_inference.py
index 61413a2..cabb02d 100644
--- a/element_facemap/facemap_inference.py
+++ b/element_facemap/facemap_inference.py
@@ -324,7 +324,10 @@ class BodyPartPosition(dj.Part):
         """
 
     def make(self, key):
-        """.populate() method will launch training for each FacemapInferenceTask"""
+        """
+        Calls facemap.pose.Pose to run pose estimation on the video files using the specified model.
+        Video files and model are specified in the FacemapInferenceTask table.
+        """
         # ID model and directories
         task_mode, output_dir = (FacemapInferenceTask & key).fetch1(
             "task_mode", "facemap_inference_output_dir"
@@ -353,27 +356,9 @@ def make(self, key):
         full_metadata_path = output_dir / f"{vid_name}_FacemapPose_metadata.pkl"
 
         # Load or Trigger Facemap Pose Estimation Inference
-        if (
+        if task_mode == "trigger" and not (
             facemap_result_path.exists() & full_metadata_path.exists()
-        ) or task_mode == "load":  # Load results and do not rerun processing
-            (
-                body_part_position_entry,
-                inference_duration,
-                total_frame_count,
-                creation_time,
-            ) = _load_facemap_results(key, facemap_result_path, full_metadata_path)
-            self.insert1(
-                {
-                    **key,
-                    "inference_completion_time": creation_time,
-                    "inference_run_duration": inference_duration,
-                    "total_frame_count": total_frame_count,
-                }
-            )
-            self.BodyPartPosition.insert(body_part_position_entry)
-            return
-
-        elif task_mode == "trigger":
+        ):
             from facemap.pose import pose as facemap_pose, model_loader
 
             bbox = (FacemapInferenceTask & key).fetch1("bbox") or []
@@ -382,9 +367,10 @@ def make(self, key):
             facemap_model_name = (
                 FacemapModel.File & f'model_id="{key["model_id"]}"'
             ).fetch1("model_file")
-
             facemap_model_path = Path.cwd() / facemap_model_name
+            # copy this model file to the facemap model root directory (~/.facemap/models/)
             models_root_dir = model_loader.get_models_dir()
+            shutil.copy(facemap_model_path, models_root_dir)
 
             # Create Symbolic Links to raw video data files from outbox directory
             video_symlinks = []
@@ -395,9 +381,6 @@ def make(self, key):
                 video_symlink.symlink_to(video_file)
                 video_symlinks.append(video_symlink.as_posix())
 
-            # copy this model file to the facemap model root directory (~/.facemap/models/)
-            shutil.copy(facemap_model_path, models_root_dir)
-
             # Instantiate Pose object, with filenames specified as video files, and bounding specified in params
             # Assumes GUI to be none as we are running CLI implementation
             pose = facemap_pose.Pose(
@@ -408,21 +391,21 @@ def make(self, key):
             )
             pose.run()
 
-            (
-                body_part_position_entry,
-                inference_duration,
-                total_frame_count,
-                creation_time,
-            ) = _load_facemap_results(key, facemap_result_path, full_metadata_path)
-            self.insert1(
-                {
-                    **key,
-                    "inference_completion_time": creation_time,
-                    "inference_run_duration": inference_duration,
-                    "total_frame_count": total_frame_count,
-                }
-            )
-            self.BodyPartPosition.insert(body_part_position_entry)
+        (
+            body_part_position_entry,
+            inference_duration,
+            total_frame_count,
+            creation_time,
+        ) = _load_facemap_results(key, facemap_result_path, full_metadata_path)
+        self.insert1(
+            {
+                **key,
+                "inference_completion_time": creation_time,
+                "inference_run_duration": inference_duration,
+                "total_frame_count": total_frame_count,
+            }
+        )
+        self.BodyPartPosition.insert(body_part_position_entry)
 
     @classmethod
     def get_trajectory(cls, key: dict, body_parts: list = "all") -> pd.DataFrame:
@@ -468,7 +451,6 @@ def get_trajectory(cls, key: dict, body_parts: list = "all") -> pd.DataFrame:
 
 def _load_facemap_results(key, facemap_result_path, full_metadata_path):
     """Load facemap results from h5 and metadata files."""
-
     from facemap import utils
 
     with open(full_metadata_path, "rb") as f: