definitively going somewhere.

dodgeballtracker/logic.py

@@ -4,7 +4,8 @@ import logging
 import torch
 from ultralytics import YOLO
 import cv2
-from dodgeballtracker.helper import create_video_writer, iter_frames
+from dodgeballtracker.helper import iter_frames
+from dodgeballtracker.video_writer import create_video_writer, VideoWriter
 import matplotlib
 import numpy as np
 from boxmot import DeepOcSort as DeepOCSORT
@@ -15,15 +16,15 @@ logger = logging.getLogger(__name__)
 Color = tuple[int, int, int]
 
 
-def get_best_available_torch_device():
+def get_best_available_torch_device() -> torch.device:
     if torch.cuda.is_available():
         logger.debug("Using the cuda (gpu) torch device.")
-        return 'cuda'
+        return torch.device('cuda:0')
     if torch.backends.mps.is_available():
         logger.debug("Using the mps (apple silicon) torch device.")
-        return 'mps'
+        return torch.device('mps')
     logger.debug("Default to cpu (no apple silicon or cuda) torch device.")
-    return 'cpu'
+    return torch.device('cpu')
 
 
 class DetectedObjectBox:
@@ -66,16 +67,21 @@ def get_court_perspective_transform():
     return cv2.getPerspectiveTransform(empirical_source, empirical_target)
 
 
+
+
+
 class DodgeballTracker:
 
     def __init__(self,
                  input_file: Path,
                  output_file: Path,
+                 video_writer: VideoWriter,
                  detected_box_objects: list[DetectedObjectBox],
                  transformation_matrix):
         self.input_path = input_file
         self.output_path = output_file
         self.output_json_path = output_file.with_suffix(".json")
+        self.video_writer = video_writer
         self.detected_box_objects = sorted(detected_box_objects, reverse=True)
         self.transformation_matrix = transformation_matrix
         self.points = []
@@ -138,7 +144,7 @@ class DodgeballTracker:
     def _get_tracker(self):
         return DeepOCSORT(
             Path('lmbn_n_cuhk03_d.pt'),  # which ReID model to use
-            torch.device(get_best_available_torch_device()),
+            get_best_available_torch_device(), # torch.cuda.get_device_name(0), #torch.device(get_best_available_torch_device()),
             True,
         )
 
@@ -150,10 +156,8 @@ class DodgeballTracker:
 
     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))
+        # writer = create_video_writer(video_cap, self.output_path)
         tracker = self._get_tracker()
 
         # load the pre-trained YOLOv8n model
@@ -168,19 +172,13 @@ class DodgeballTracker:
                 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)
 
-            self.write_frame(writer, frame)
+            self.video_writer.write_frame(frame)
 
             if cv2.waitKey(1) == ord("q"):
                 break
-        print("fini")
 
         video_cap.release()
-        writer.release()
-        cv2.destroyAllWindows()
         self.persist_points()
-        print("finished")