dodgeballtracker/helper.py

@@ -1,5 +1,7 @@
 import contextlib
 import time
+from typing import Iterator
+from dodgeballtracker.typing import VideoFrame, VideoTimestamp
 import cv2
 import logging
 
@@ -7,7 +9,7 @@ logger = logging.getLogger(__name__)
 
 
 
-def iter_frames(video_cap):
+def iter_frames(video_cap) -> Iterator[tuple[VideoFrame, VideoTimestamp]]:
     while 1:
         ret, frame = video_cap.read()
         if not ret:

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()

dodgeballtracker/typing.py

@@ -0,0 +1,7 @@
+from typing import NewType
+
+type XCoord = float
+type YCoord = NewType('YCoord', float)
+type Point = tuple[XCoord, YCoord]
+type VideoFrame = NewType('Frame', float)
+type VideoTimestamp = float

dodgeballtracker/video_writer.py

@@ -1,6 +1,6 @@
 import abc
 from pathlib import Path
-
+from dodgeballtracker.typing import VideoFrame
 import cv2
 
 def create_video_writer(video_cap: cv2.VideoCapture, output_filename: Path) -> cv2.VideoWriter:
@@ -23,35 +23,45 @@ def create_video_writer(video_cap: cv2.VideoCapture, output_filename: Path) -> c
 class VideoWriter(abc.ABC):
 
     @abc.abstractmethod
-    def write_frame(self, frame) -> None:
+    def write_frame(self, frame: VideoFrame) -> None:
         pass
 
     @abc.abstractmethod
     def close(self) -> None:
         pass
 
+    @abc.abstractmethod
+    def wait(self) -> bool:
+        pass
+
 
 class NoopVideoWriter(VideoWriter):
 
-    def write_frame(self, frame) -> None:
+    def write_frame(self, frame: VideoFrame) -> None:
         pass
 
     def close(self) -> None:
         pass
 
+    def wait(self) -> bool:
+        return False
+
 
 class MP4VVideoWriter(VideoWriter):
 
     def __init__(self, video_writer: cv2.VideoWriter) -> None:
         self._video_writer = video_writer
 
-    def write_frame(self, frame) -> None:
+    def write_frame(self, frame: VideoFrame) -> None:
         self._video_writer.write(frame)
 
     def close(self) -> None:
         self._video_writer.release()
         cv2.destroyAllWindows()
 
+    def wait(self) -> bool:
+        return False #cv2.waitKey(1) == ord("q")
+
     @classmethod
     def from_frame(cls, input_: Path, output: Path) -> VideoWriter:
         video_cap = cv2.VideoCapture(str(input_))
@@ -62,6 +72,9 @@ class MP4VVideoWriter(VideoWriter):
 
 
 class MP4VVideoWriterWithFrame(MP4VVideoWriter):
-    def write_frame(self, frame) -> None:
+    def write_frame(self, frame: VideoFrame) -> None:
         super().write_frame(frame)
         cv2.imshow("Frame", frame)
+
+    def wait(self) -> bool:
+        return cv2.waitKey(1) == ord("q")