cleaned up code.

dodgeballtracker/__main__.py

@@ -41,12 +41,7 @@ def get_video_writer(display_frame: bool,
     return closing(MP4VVideoWriter.from_frame(video_input, video_output, thresholds))
 
 
-@cli.command()
-@click.option('--input', '-i', 'input_',
-              type=click.Path(readable=True, path_type=Path))
-@click.option('--output', '-o', type=click.Path(path_type=Path))
-@click.option('--display-video', '-d', type=bool,  is_flag=True, show_default=True, default=False)
-def process(input_: Path, output: Path | None, display_video: bool):
+def process_logic(input_: Path, output: Path | None, display_video: bool):
     output = get_output(input_, output)
     detected_box_objects = [
         DetectedObjectBox.from_viridis(threshold) for threshold in
@@ -65,6 +60,15 @@ def process(input_: Path, output: Path | None, display_video: bool):
         with time_tracker():
             tracker.main()
 
+@cli.command()
+@click.option('--input', '-i', 'input_',
+              type=click.Path(readable=True, path_type=Path))
+@click.option('--output', '-o', type=click.Path(path_type=Path))
+@click.option('--display-video', '-d', type=bool,  is_flag=True, show_default=True, default=False)
+def process(input_: Path, output: Path | None, display_video: bool):
+    process_logic(input_, output, display_video)
+
 
 if __name__ == "__main__":
     cli()
+    # process_logic(Path('/home/didier/all-dressed-programming/dodgeball-tracker/data/micro.mp4'), None, False)

dodgeballtracker/logic.py

@@ -10,7 +10,7 @@ import numpy as np
 import pandas as pd
 import pandera as pa
 from boxmot import DeepOcSort as DeepOCSORT
-from .typing import VideoFrame, DetectedObject_
+from .typing import VideoFrame, DetectedObject_, TrackedDetectedObject_
 
 logger = logging.getLogger(__name__)
 
@@ -111,7 +111,7 @@ class DodgeballTracker:
     def run_model_on_frame(self, model, frame: VideoFrame):
         model_outputs = model(frame)[0]
         boxes = model_outputs.boxes.data.cpu().numpy()
-        columns=['xmin', 'ymin', 'xmax', 'ymin', 'score', 'category']
+        columns=['xmin', 'ymin', 'xmax', 'ymax', 'score', 'category']
         return pd.DataFrame(boxes, columns=columns)
 
     def persist_points(self):
@@ -131,9 +131,10 @@ class DodgeballTracker:
         """
         return YOLO("yolo11n.pt", verbose=True)
 
-    def update_tracker(self, tracker: DeepOCSORT, detected_objs, frame: VideoFrame):
+    @pa.check_io(detected_objs=DetectedObject_, out=TrackedDetectedObject_)
+    def update_tracker(self, tracker: DeepOCSORT, detected_objs: pd.DataFrame, frame: VideoFrame):
         data_ = tracker.update(detected_objs.to_numpy(), frame)
-        columns = ['xmin', 'ymin', 'xmax', 'ymin', 'object_id', 'score', 'category', 'index']
+        columns = ['xmin', 'ymin', 'xmax', 'ymax', 'object_id', 'score', 'category', 'index']
         return pd.DataFrame(data_, columns=columns)
 
 
@@ -147,9 +148,8 @@ class DodgeballTracker:
             # run the YOLO model on the frame
             detected_objs = self.run_model_on_frame(model, frame)
             tracked_obj = self.update_tracker(tracker, detected_objs, frame)
-            # self.video_writer.plot_boxes_to_frame(detected_objs, frame)
-            #
-            # self.video_writer.write_frame(frame)
+            self.video_writer.plot_boxes_to_frame(detected_objs, frame)
+            self.video_writer.write_frame(frame)
 
             if self.video_writer.wait():
                 break

dodgeballtracker/typing.py

@@ -2,7 +2,7 @@ from typing import NewType
 from dataclasses import dataclass
 import numpy as np
 import pandera as pa
-from pandera.dtypes import Float32
+from pandera.dtypes import Float32, Float64
 
 XCoord = NewType('XCoord', float)
 YCoord = NewType('YCoord', float)
@@ -10,15 +10,16 @@ Score = 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
@@ -29,31 +30,24 @@ class DetectedObject:
     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))
+        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.score, self.category])
+        return np.array([self.xmin, self.ymax, self.xmax, self.ymax, self.score,
+                         self.category])
 
 
-DetectedObject_ = pa.DataFrameSchema({
-    "xmin": pa.Column(Float32),
-    "ymin": pa.Column(Float32),
-    "xmax": pa.Column(Float32),
-    "ymin": pa.Column(Float32),
-    "score": pa.Column(Float32),
-    "category": pa.Column(Float32),
-})
+def tensor_of(data_type: pa.DataType, cols: list[str]) -> pa.DataFrameSchema:
+    return pa.DataFrameSchema({col: pa.Column(data_type) for col in cols})
 
-TrackedDetectedObject_ = pa.DataFrameSchema({
-    "xmin": pa.Column(Float32),
-    "ymin": pa.Column(Float32),
-    "xmax": pa.Column(Float32),
-    "ymin": pa.Column(Float32),
-    "object_id": pa.Column(Float32),
-    "score": pa.Column(Float32),
-    "category": pa.Column(Float32),
-    "index": pa.Column(Float32),
-})
+
+DetectedObject_ = tensor_of(Float32, ["xmin", "ymin", "xmax", "ymax", "score",
+                                      "category"])
+
+TrackedDetectedObject_ = tensor_of(Float64,
+                                   list(DetectedObject_.columns) + ["object_id", "index"])
+# ["xmin", "ymin", "xmax", "ymax", "object_id", "score", "category", "index"])
 
 
 Color = tuple[int, int, int]

dodgeballtracker/video_writer.py

@@ -1,6 +1,10 @@
 import abc
 from pathlib import Path
-from dodgeballtracker.typing import VideoFrame, DetectedObject, Color
+
+import pandas as pd
+import pandera as pa
+
+from dodgeballtracker.typing import VideoFrame, DetectedObject, Color, DetectedObject_
 import matplotlib
 import cv2
 
@@ -44,7 +48,7 @@ class VideoWriter(abc.ABC):
         pass
 
     @abc.abstractmethod
-    def plot_boxes_to_frame(self, data: list[DetectedObject], frame: VideoFrame) -> None:
+    def plot_boxes_to_frame(self, data: pd.DataFrame, frame: VideoFrame) -> None:
         pass
 
     @abc.abstractmethod
@@ -58,7 +62,7 @@ class VideoWriter(abc.ABC):
 
 class NoopVideoWriter(VideoWriter):
 
-    def plot_boxes_to_frame(self, boxes: list[DetectedObject], frame: VideoFrame) -> None:
+    def plot_boxes_to_frame(self, boxes: pd.DataFrame, frame: VideoFrame) -> None:
         pass
 
     def write_frame(self, frame: VideoFrame) -> None:
@@ -80,16 +84,20 @@ class MP4VVideoWriter(VideoWriter):
     def write_frame(self, frame: VideoFrame) -> None:
         self._video_writer.write(frame)
 
-    def plot_box_to_frame(self, data: DetectedObject, frame: VideoFrame) -> None:
+    def plot_box_to_frame(self, data: pd.Series, frame: VideoFrame) -> None:
         confidence = float(data.score)
+        # data = data.map(int)
         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)
+                start_point =  data[['xmin', 'ymin']].array.astype('int32')
+                end_point = data[['xmax', 'ymax']].array.astype('int32')
+                cv2.rectangle(frame, start_point, end_point, threshold.color, 2)
                 return
 
-    def plot_boxes_to_frame(self, boxes: list[DetectedObject], frame: VideoFrame) -> None:
-        for box in boxes:
+    @pa.decorators.check_io(boxes=DetectedObject_)
+    def plot_boxes_to_frame(self, boxes: pd.DataFrame, frame: VideoFrame) -> None:
+        boxes_that_passe_threshold = boxes[boxes['score'] > 0.5]
+        for i, box in boxes_that_passe_threshold.iterrows():
             self.plot_box_to_frame(box, frame)