cleaned up code.

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