added pandera

dodgeballtracker/logic.py

@@ -7,8 +7,10 @@ import cv2
 from dodgeballtracker.helper import iter_frames
 from dodgeballtracker.video_writer import VideoWriter, DetectedObjectBox
 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_
 
 logger = logging.getLogger(__name__)
 
@@ -105,21 +107,18 @@ class DodgeballTracker:
                 return True
         return False
 
+    @pa.check_output(DetectedObject_)
     def run_model_on_frame(self, model, frame: VideoFrame):
         model_outputs = model(frame)[0]
-        return model_outputs.boxes.data.cpu().numpy()
-        # return model_outputs.boxes.data.tolist()
-        foo = model_outputs.boxes.data
-        for data in model_outputs.boxes.data.tolist():
-            detected_obj = DetectedObject(*data)
-            if detected_obj.score > 0.5:
-                yield detected_obj
+        boxes = model_outputs.boxes.data.cpu().numpy()
+        columns=['xmin', 'ymin', 'xmax', 'ymin', 'score', 'category']
+        return pd.DataFrame(boxes, columns=columns)
 
     def persist_points(self):
         with open(self.output_json_path, "wt") as file_:
             json.dump({"points": self.points}, file_, indent=4)
 
-    def _get_tracker(self):
+    def _get_tracker(self) -> DeepOCSORT:
         return DeepOCSORT(
             Path('lmbn_n_cuhk03_d.pt'),  # which ReID model to use
             get_best_available_torch_device(),
@@ -132,8 +131,11 @@ class DodgeballTracker:
         """
         return YOLO("yolo11n.pt", verbose=True)
 
-    def update_tracker(self, tracker, detected_objs: list[DetectedObject], frame: VideoFrame) -> None:
-        important_detected_objs = [do for do in detected_objs if do.score > 0.5]
+    def update_tracker(self, tracker: DeepOCSORT, detected_objs, frame: VideoFrame):
+        data_ = tracker.update(detected_objs.to_numpy(), frame)
+        columns = ['xmin', 'ymin', 'xmax', 'ymin', 'object_id', 'score', 'category', 'index']
+        return pd.DataFrame(data_, columns=columns)
+
 
     def main(self):
         # define some constants
@@ -143,18 +145,11 @@ class DodgeballTracker:
 
         for i, (frame, ts) in enumerate(iter_frames(video_cap)):
             # run the YOLO model on the frame
-            detected_objs = list(self.run_model_on_frame(model, frame))
-            data_ = tracker.update(
-                np.array(detected_objs), frame)
-            if i == 100:
-                print(data_)
+            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)
-            # data_ = tracker.update(np.array([dobj.as_nparray() for dobj in detected_objs]), frame)
-            # if data_:
-            #     print("dkld")
-            # self.print_tracker_data(data_, i, ts)
-
-            self.video_writer.write_frame(frame)
+            #
+            # self.video_writer.write_frame(frame)
 
             if self.video_writer.wait():
                 break

dodgeballtracker/typing.py

@@ -1,6 +1,8 @@
 from typing import NewType
 from dataclasses import dataclass
 import numpy as np
+import pandera as pa
+from pandera.dtypes import Float32
 
 XCoord = NewType('XCoord', float)
 YCoord = NewType('YCoord', float)
@@ -33,5 +35,25 @@ class DetectedObject:
         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),
+})
+
+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),
+})
+
 
 Color = tuple[int, int, int]