Commiting with problems in tracking.

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