changed tracker

dodgeballtracker/logic.py

@@ -1,12 +1,14 @@
 import json
 from pathlib import Path
 
+import torch
 from ultralytics import YOLO
 import cv2
 from dodgeballtracker.helper import create_video_writer, iter_frames
-from deep_sort_realtime.deepsort_tracker import DeepSort
 import matplotlib
 import numpy as np
+from boxmot import DeepOCSORT
+
 
 
 Color = tuple[int, int, int]
@@ -87,7 +89,7 @@ class DodgeballTracker:
                 xmin, ymin, xmax, ymax = self.extract_box(data)
                 self.print_box(data, frame_idx, frame_ts)
                 cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2)
-                return [[xmin, ymin, xmax - xmin, ymax - ymin], confidence, int(data[5])]
+                return data  # [[xmin, ymin, xmax - xmin, ymax - ymin], confidence, int(data[5])]
 
     def run_model_on_frame(self, model, frame):
         return model(frame)[0]
@@ -114,7 +116,11 @@ class DodgeballTracker:
 
         video_cap = cv2.VideoCapture(str(self.input_path))
         writer = create_video_writer(video_cap, str(self.output_path))
-        tracker = DeepSort(max_age=50)
+        tracker = DeepOCSORT(
+            model_weights=Path('osnet_x0_25_msmt17.pt'),  # which ReID model to use
+            device=torch.device("mps"),
+            fp16=True,
+        )
 
         # load the pre-trained YOLOv8n model
         model = YOLO("yolov8n.pt")
@@ -126,14 +132,7 @@ class DodgeballTracker:
             for data in detections.boxes.data.tolist():
                 results.append(self.plot_boxe_to_frame(data, frame, i, ts))
 
-            tracks = tracker.update_tracks([r for r in results if r is not None], frame=frame)
-            for track in tracks:
-                if not track.is_confirmed():
-                    continue
-                track_id = track.track_id
-                ltrb = track.to_ltrb()
-                print(f"{track_id} => { ltrb }")
-
+            tracks = tracker.update(np.array([r for r in results if r is not None]), img=frame)
 
             # show the frame to our screen
             self.write_frame(writer, frame)