cleaned up code.

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)