expose model info in __init__

freemocap · Jun 18, 2024 · 1e743c3 · 1e743c3
1 parent c480f03
commit 1e743c3
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/skellytracker/__init__.py b/skellytracker/__init__.py
@@ -27,10 +27,12 @@
 
 try:
     from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import MediapipeHolisticTracker
+    from skellytracker.trackers.mediapipe_tracker.mediapipe_model_info import MediapipeModelInfo
 except:
     print("To use mediapipe_holistic_tracker, install skellytracker[mediapipe]")
 try:
     from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker
+    from skellytracker.trackers.yolo_tracker.yolo_model_info import YOLOModelInfo
 except:
     print("To use yolo_tracker, install skellytracker[yolo]")
 try:

diff --git a/skellytracker/trackers/base_tracker/base_recorder.py b/skellytracker/trackers/base_tracker/base_recorder.py
@@ -53,6 +53,8 @@ def save(self, file_path: str) -> None:
         :return: None
         """
         if self.recorded_objects_array is None:
-            self.process_tracked_objects()
+            recorded_objects_array = self.process_tracked_objects()
+        else:
+            recorded_objects_array = self.recorded_objects_array
         logger.info(f"Saving recorded objects to {file_path}")
-        np.save(file_path, self.recorded_objects_array)
+        np.save(file_path, recorded_objects_array)