Migrated to nix for dependency managment

dodgeballtracker/__init__.py

@@ -1,24 +0,0 @@
-from pathlib import Path
-
-from dodgeballtracker.logic import (DetectedObjectBox,
-                                    DodgeballTracker,
-                                    get_court_perspective_transform)
-import numpy as np
-
-
-here = Path(__file__).parent
-input_dir = here / '..' / 'data'
-output_dir = here / '..' / 'outputs'
-
-if __name__ == "__main__":
-    detected_box_objects = [
-        DetectedObjectBox.from_viridis(threshold) for threshold in np.arange(0.5, 1, 0.05)
-    ]
-    input_video = input_dir / "micro.mp4"
-    output_video = output_dir / "output-lmbn_n_cuhk03_d.mp4"
-
-    transformation_matrix = get_court_perspective_transform()
-
-    dodgeball_tracker = DodgeballTracker(input_video, output_video, detected_box_objects, transformation_matrix)
-
-    dodgeball_tracker.main()

dodgeballtracker/__main__.py

@@ -0,0 +1,28 @@
+from pathlib import Path
+
+from dodgeballtracker.logic import (DetectedObjectBox,
+                                    DodgeballTracker,
+                                    get_court_perspective_transform)
+import numpy as np
+import logging
+
+
+here = Path(__file__).parent
+input_dir = here / '..' / 'data'
+output_dir = here / '..' / 'outputs'
+
+if __name__ == "__main__":
+    logging.basicConfig(filename='tracker.log', level=logging.DEBUG)
+
+    detected_box_objects = [
+        DetectedObjectBox.from_viridis(threshold) for threshold in np.arange(0.5, 1, 0.05)
+    ]
+
+    input_video = input_dir / "small.mp4"
+    output_video = output_dir / "foo.mp4"
+
+    transformation_matrix = get_court_perspective_transform()
+
+    dodgeball_tracker = DodgeballTracker(input_video, output_video, detected_box_objects, transformation_matrix)
+
+    dodgeball_tracker.main()

dodgeballtracker/logic.py

@@ -7,7 +7,7 @@ import cv2
 from dodgeballtracker.helper import create_video_writer, iter_frames
 import matplotlib
 import numpy as np
-from boxmot import DeepOCSORT
+from boxmot import DeepOcSort as DeepOCSORT
 
 
 
@@ -125,19 +125,20 @@ class DodgeballTracker:
 
     def _get_tracker(self):
         return DeepOCSORT(
-            model_weights=Path('lmbn_n_cuhk03_d.pt'),  # which ReID model to use
-            device=torch.device("mps"),
-            fp16=True,
+            Path('lmbn_n_cuhk03_d.pt'),  # which ReID model to use
+            torch.device("cpu"),
+            True,
         )
 
     def _get_detection_model(self):
         """
         load the pre-trained YOLOv8n model
         """
-        return YOLO("yolov8n.pt")
+        return YOLO("yolo11n.pt", verbose=True)
 
     def main(self):
         # define some constants
+        print("started")
 
         video_cap = cv2.VideoCapture(str(self.input_path))
         writer = create_video_writer(video_cap, str(self.output_path))
@@ -150,12 +151,15 @@ class DodgeballTracker:
             # run the YOLO model on the frame
             res = []
             detections = self.run_model_on_frame(model, frame)
+            print("ici?")
             for data in detections.boxes.data.tolist():
+                print("la?")
                 is_good_box = self.plot_boxe_to_frame(data, frame)
                 if is_good_box:
                     res.append(data)
 
             # show the frame to our screen
+            print("tracker?")
             data_ = tracker.update(np.array(res), frame)
             self.print_tracker_data(data_, i, ts)
 
@@ -163,8 +167,10 @@ class DodgeballTracker:
 
             if cv2.waitKey(1) == ord("q"):
                 break
+        print("fini")
 
         video_cap.release()
         writer.release()
         cv2.destroyAllWindows()
         self.persist_points()
+        print("finished")