definitively going somewhere.

dodgeballtracker/logic.py

@@ -1,4 +1,5 @@
 import json
+from collections import Counter
 from pathlib import Path
 import logging
 import torch
@@ -9,7 +10,7 @@ from dodgeballtracker.video_writer import create_video_writer, VideoWriter
 import matplotlib
 import numpy as np
 from boxmot import DeepOcSort as DeepOCSORT
-
+from .typing import VideoFrame
 
 logger = logging.getLogger(__name__)
 
@@ -111,7 +112,7 @@ class DodgeballTracker:
         for d in tracker_data:
             self.print_box(d, frame_idx, frame_ts)
 
-    def plot_box_to_frame(self, data, frame):
+    def plot_box_to_frame(self, data, frame: VideoFrame):
         confidence = float(data[4])
         for detected_box_object in self.detected_box_objects:
 
@@ -121,22 +122,9 @@ class DodgeballTracker:
                 return True
         return False
 
-    def run_model_on_frame(self, model, frame):
+    def run_model_on_frame(self, model, frame: VideoFrame):
         return model(frame)[0]
 
-    def crazy_experiment(self, frame):
-        cv2.putText(frame,
-                    "Hello world",
-                    (50, 50),
-                    cv2.FONT_HERSHEY_SIMPLEX,
-                    2,
-                    (0, 255, 0), 8)
-        cv2.line(frame, (0, 0), (511, 511), (255, 0, 0), 5)
-        pts = np.array([[50, 50], [500, 50], [500, 500], [50, 500], [50, 50]], np.int32)
-        pts = pts.reshape((-1, 1, 2))
-
-        cv2.polylines(frame, pts, True, (0, 255, 255), 8)
-
     def persist_points(self):
         with open(self.output_json_path, "wt") as file_:
             json.dump({"points": self.points}, file_, indent=4)
@@ -144,7 +132,7 @@ class DodgeballTracker:
     def _get_tracker(self):
         return DeepOCSORT(
             Path('lmbn_n_cuhk03_d.pt'),  # which ReID model to use
-            get_best_available_torch_device(), # torch.cuda.get_device_name(0), #torch.device(get_best_available_torch_device()),
+            get_best_available_torch_device(),
             True,
         )
 
@@ -157,11 +145,9 @@ class DodgeballTracker:
     def main(self):
         # define some constants
         video_cap = cv2.VideoCapture(str(self.input_path))
-        # writer = create_video_writer(video_cap, self.output_path)
         tracker = self._get_tracker()
-
-        # load the pre-trained YOLOv8n model
         model = self._get_detection_model()
+        counter = Counter()
 
         for i, (frame, ts) in enumerate(iter_frames(video_cap)):
             # run the YOLO model on the frame
@@ -177,8 +163,9 @@ class DodgeballTracker:
 
             self.video_writer.write_frame(frame)
 
-            if cv2.waitKey(1) == ord("q"):
+            if self.video_writer.wait():
                 break
 
+        print(f"counter: {counter}")
         video_cap.release()
         self.persist_points()