Going somewhere without a glass of wine.

dodgeballtracker/video_writer.py

@@ -4,7 +4,7 @@ from pathlib import Path
 import pandas as pd
 import pandera as pa
 
-from dodgeballtracker.typing import VideoFrame, Color, DetectedObject_
+from dodgeballtracker.typing import (VideoFrame, Color, TrackedDetectedObject)
 import matplotlib
 import cv2
 
@@ -84,17 +84,34 @@ class MP4VVideoWriter(VideoWriter):
     def write_frame(self, frame: VideoFrame) -> None:
         self._video_writer.write(frame)
 
+    def write_box_id(self, frame: VideoFrame, box: pd.Series):
+        bottom_left_corner_of_text = (int(box['xmin']), int(box['ymax']))
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        font_scale = 0.5
+        font_color = (255, 255, 255)
+        thickness = 1
+        line_type = 2
+        id_ = str(int(box['object_id']))
+        cv2.putText(frame,
+                    id_,
+                    bottom_left_corner_of_text,
+                    font,
+                    font_scale,
+                    font_color,
+                    thickness,
+                    line_type)
+
     def plot_box_to_frame(self, data: pd.Series, frame: VideoFrame) -> None:
-        confidence = float(data.score)
-        # data = data.map(int)
+        confidence = data['score']
         for threshold in self._thresholds:
             if threshold.min_threshold < confidence:
                 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)
+                self.write_box_id(frame, data)
                 return
 
-    @pa.decorators.check_io(boxes=DetectedObject_)
+    @pa.decorators.check_io(boxes=TrackedDetectedObject)
     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():