dodgeballtracker/__main__.py

@@ -1,17 +1,17 @@
-import time
 from contextlib import closing
 from pathlib import Path
 
 from dodgeballtracker.helper import time_tracker
-from dodgeballtracker.logic import (DetectedObjectBox,
-                                    DodgeballTracker,
+from dodgeballtracker.logic import (DodgeballTracker,
                                     get_court_perspective_transform)
 import numpy as np
 import logging
 import click
 from typing import ContextManager
 
-from dodgeballtracker.video_writer import (MP4VVideoWriter, VideoWriter,
+from dodgeballtracker.video_writer import (MP4VVideoWriter,
+                                           VideoWriter,
+                                           DetectedObjectBox,
                                            MP4VVideoWriterWithFrame)
 
 here = Path(__file__).parent
@@ -34,10 +34,11 @@ def get_output(input_: Path, proposed_output: Path | None) -> Path:
 def get_video_writer(display_frame: bool,
                      video_input: Path,
                      video_output: Path) -> ContextManager[VideoWriter]:
+    thresholds = list(np.arange(0.5, 1, 0.05))
     if display_frame:
         return closing(
-            MP4VVideoWriterWithFrame.from_frame(video_input, video_output))
-    return closing(MP4VVideoWriter.from_frame(video_input, video_output))
+            MP4VVideoWriterWithFrame.from_frame(video_input, video_output, thresholds))
+    return closing(MP4VVideoWriter.from_frame(video_input, video_output, thresholds))
 
 
 @cli.command()

dodgeballtracker/logic.py

@@ -1,16 +1,14 @@
 import json
-from collections import Counter
 from pathlib import Path
 import logging
 import torch
 from ultralytics import YOLO
 import cv2
 from dodgeballtracker.helper import iter_frames
-from dodgeballtracker.video_writer import create_video_writer, VideoWriter
-import matplotlib
+from dodgeballtracker.video_writer import VideoWriter, DetectedObjectBox
 import numpy as np
 from boxmot import DeepOcSort as DeepOCSORT
-from .typing import VideoFrame
+from .typing import VideoFrame, DetectedObject
 
 logger = logging.getLogger(__name__)
 
@@ -28,21 +26,6 @@ def get_best_available_torch_device() -> torch.device:
     return torch.device('cpu')
 
 
-class DetectedObjectBox:
-    def __init__(self, min_threshold: float, color: Color):
-        self.min_threshold = min_threshold
-        self.color = color
-
-    def __lt__(self, other):
-        return self.min_threshold < other.min_threshold
-
-    @classmethod
-    def from_viridis(cls, min_threshold):
-        cm = matplotlib.colormaps['viridis']
-        color: Color = tuple((int(255 * cm(min_threshold)[i]) for i in range(3)))
-        return cls(min_threshold, color)
-
-
 def get_court_perspective_transform():
     """
     """
@@ -117,13 +100,17 @@ class DodgeballTracker:
         for detected_box_object in self.detected_box_objects:
 
             if detected_box_object.min_threshold < confidence:
-                xmin, ymin, xmax, ymax = self.extract_box(data)
-                cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2)
+                # xmin, ymin, xmax, ymax = self.extract_box(data)
+                # cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2)
                 return True
         return False
 
     def run_model_on_frame(self, model, frame: VideoFrame):
-        return model(frame)[0]
+        model_outputs = model(frame)[0]
+        for data in model_outputs.boxes.data.tolist():
+            detected_obj = DetectedObject(*data)
+            if detected_obj.confidence > 0.5:
+                yield detected_obj
 
     def persist_points(self):
         with open(self.output_json_path, "wt") as file_:
@@ -142,30 +129,28 @@ class DodgeballTracker:
         """
         return YOLO("yolo11n.pt", verbose=True)
 
+    def update_tracker(self, tracker, detected_objs: list[DetectedObject], frame: VideoFrame) -> None:
+        important_detected_objs = [do for do in detected_objs if do.confidence > 0.5]
+
     def main(self):
         # define some constants
         video_cap = cv2.VideoCapture(str(self.input_path))
         tracker = self._get_tracker()
         model = self._get_detection_model()
-        counter = Counter()
 
         for i, (frame, ts) in enumerate(iter_frames(video_cap)):
             # run the YOLO model on the frame
-            res = []
-            detections = self.run_model_on_frame(model, frame)
-            for data in detections.boxes.data.tolist():
-                is_good_box = self.plot_box_to_frame(data, frame)
-                if is_good_box:
-                    res.append(data)
-
-            data_ = tracker.update(np.array(res), frame)
-            self.print_tracker_data(data_, i, ts)
+            detected_objs = list(self.run_model_on_frame(model, frame))
+            self.video_writer.plot_boxes_to_frame(detected_objs, frame)
+            # data_ = tracker.update(np.array([dobj.as_nparray() for dobj in detected_objs]), frame)
+            # if data_:
+            #     print("dkld")
+            # self.print_tracker_data(data_, i, ts)
 
             self.video_writer.write_frame(frame)
 
             if self.video_writer.wait():
                 break
 
-        print(f"counter: {counter}")
         video_cap.release()
         self.persist_points()

dodgeballtracker/typing.py

@@ -1,7 +1,37 @@
 from typing import NewType
+from dataclasses import dataclass
+import numpy as np
 
-type XCoord = float
-type YCoord = NewType('YCoord', float)
-type Point = tuple[XCoord, YCoord]
-type VideoFrame = NewType('Frame', float)
-type VideoTimestamp = float
+XCoord = NewType('XCoord', float)
+YCoord = NewType('YCoord', float)
+Confidence = NewType('Confidence', float)
+Category = NewType('Category', float)
+
+
+
+@dataclass
+class Point:
+    x: XCoord
+    y: YCoord
+
+VideoFrame = NewType('Frame', float)
+VideoTimestamp = float
+
+@dataclass
+class DetectedObject:
+    xmin: XCoord
+    ymin: YCoord
+    xmax: XCoord
+    ymax: YCoord
+    confidence: Confidence
+    category: Category
+
+    def as_int_box(self) -> tuple[tuple[int, int], tuple[int, int]]:
+        return (int(self.xmin), int(self.ymin)), (int(self.xmax), int(self.ymax))
+
+    def as_nparray(self):
+        return np.array([self.xmin, self.ymax, self.xmax, self.ymax, self.confidence, self.category])
+
+
+
+Color = tuple[int, int, int]

dodgeballtracker/video_writer.py

@@ -1,8 +1,25 @@
 import abc
 from pathlib import Path
-from dodgeballtracker.typing import VideoFrame
+from dodgeballtracker.typing import VideoFrame, DetectedObject, Color
+import matplotlib
 import cv2
 
+
+class DetectedObjectBox:
+    def __init__(self, min_threshold: float, color: Color):
+        self.min_threshold = min_threshold
+        self.color = color
+
+    def __lt__(self, other):
+        return self.min_threshold < other.min_threshold
+
+    @classmethod
+    def from_viridis(cls, min_threshold):
+        cm = matplotlib.colormaps['viridis']
+        color: Color = tuple((int(255 * cm(min_threshold)[i]) for i in range(3)))
+        return cls(min_threshold, color)
+
+
 def create_video_writer(video_cap: cv2.VideoCapture, output_filename: Path) -> cv2.VideoWriter:
 
     # grab the width, height, and fps of the frames in the video stream.
@@ -26,6 +43,10 @@ class VideoWriter(abc.ABC):
     def write_frame(self, frame: VideoFrame) -> None:
         pass
 
+    @abc.abstractmethod
+    def plot_boxes_to_frame(self, data: list[DetectedObject], frame: VideoFrame) -> None:
+        pass
+
     @abc.abstractmethod
     def close(self) -> None:
         pass
@@ -37,6 +58,9 @@ class VideoWriter(abc.ABC):
 
 class NoopVideoWriter(VideoWriter):
 
+    def plot_boxes_to_frame(self, boxes: list[DetectedObject], frame: VideoFrame) -> None:
+        pass
+
     def write_frame(self, frame: VideoFrame) -> None:
         pass
 
@@ -49,24 +73,39 @@ class NoopVideoWriter(VideoWriter):
 
 class MP4VVideoWriter(VideoWriter):
 
-    def __init__(self, video_writer: cv2.VideoWriter) -> None:
+    def __init__(self, video_writer: cv2.VideoWriter, thresholds: list[DetectedObjectBox]) -> None:
         self._video_writer = video_writer
+        self._thresholds = thresholds
 
     def write_frame(self, frame: VideoFrame) -> None:
         self._video_writer.write(frame)
 
+    def plot_box_to_frame(self, data: DetectedObject, frame: VideoFrame) -> None:
+        confidence = float(data.confidence)
+        for threshold in self._thresholds:
+            if threshold.min_threshold < confidence:
+                # xmin, ymin, xmax, ymax = self._extract_box(data)
+                cv2.rectangle(frame, *data.as_int_box(), threshold.color, 2)
+                return
+
+    def plot_boxes_to_frame(self, boxes: list[DetectedObject], frame: VideoFrame) -> None:
+        for box in boxes:
+            self.plot_box_to_frame(box, frame)
+
+
     def close(self) -> None:
         self._video_writer.release()
         cv2.destroyAllWindows()
 
     def wait(self) -> bool:
-        return False #cv2.waitKey(1) == ord("q")
+        return False
 
     @classmethod
-    def from_frame(cls, input_: Path, output: Path) -> VideoWriter:
+    def from_frame(cls, input_: Path, output: Path, thresholds: list[float]) -> VideoWriter:
+        object_drawers = [DetectedObjectBox.from_viridis(t) for t in thresholds]
         video_cap = cv2.VideoCapture(str(input_))
         try:
-            return cls(create_video_writer(video_cap, output))
+            return cls(create_video_writer(video_cap, output), object_drawers)
         finally:
             video_cap.release()
 

flake.nix

@@ -27,6 +27,7 @@
         virtual-env = pkgs.python312.withPackages (ps: [
           ps.matplotlib
           ps.torchWithCuda
+          ps.jupyter
           pythonPackages.opencv4
           pythonPackages.boxmot
           pythonPackages.ultralytics