modidied stuff

dodgeballtracker/logic.py

@@ -1,6 +1,6 @@
 import json
 from pathlib import Path
-
+import logging
 import torch
 from ultralytics import YOLO
 import cv2
@@ -10,10 +10,22 @@ import numpy as np
 from boxmot import DeepOcSort as DeepOCSORT
 
 
+logger = logging.getLogger(__name__)
 
 Color = tuple[int, int, int]
 
 
+def get_best_available_torch_device():
+    if torch.cuda.is_available():
+        logger.debug("Using the cuda (gpu) torch device.")
+        return 'cuda'
+    if torch.backends.mps.is_available():
+        logger.debug("Using the mps (apple silicon) torch device.")
+        return 'mps'
+    logger.debug("Default to cpu (no apple silicon or cuda) torch device.")
+    return 'cpu'
+
+
 class DetectedObjectBox:
     def __init__(self, min_threshold: float, color: Color):
         self.min_threshold = min_threshold
@@ -93,7 +105,7 @@ class DodgeballTracker:
         for d in tracker_data:
             self.print_box(d, frame_idx, frame_ts)
 
-    def plot_boxe_to_frame(self, data, frame):
+    def plot_box_to_frame(self, data, frame):
         confidence = float(data[4])
         for detected_box_object in self.detected_box_objects:
 
@@ -126,7 +138,7 @@ class DodgeballTracker:
     def _get_tracker(self):
         return DeepOCSORT(
             Path('lmbn_n_cuhk03_d.pt'),  # which ReID model to use
-            torch.device("cpu"),
+            torch.device(get_best_available_torch_device()),
             True,
         )
 
@@ -151,10 +163,8 @@ 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)
+                is_good_box = self.plot_box_to_frame(data, frame)
                 if is_good_box:
                     res.append(data)