Going somewhere with a glass of wine.

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_, TrackedDetectedObject_
+from .typing import VideoFrame, DetectedObject_, TrackedDetectedObject
 
 logger = logging.getLogger(__name__)
 
@@ -112,7 +112,8 @@ class DodgeballTracker:
         model_outputs = model(frame)[0]
         boxes = model_outputs.boxes.data.cpu().numpy()
         columns=['xmin', 'ymin', 'xmax', 'ymax', 'score', 'category']
-        return pd.DataFrame(boxes, columns=columns)
+        ret = pd.DataFrame(boxes, columns=columns)
+        return ret
 
     def persist_points(self):
         with open(self.output_json_path, "wt") as file_:
@@ -129,9 +130,9 @@ class DodgeballTracker:
         """
         load the pre-trained YOLOv8n model
         """
-        return YOLO("yolo11n.pt", verbose=True)
+        return YOLO("yolo11n.pt", verbose=False)
 
-    @pa.check_io(detected_objs=DetectedObject_, out=TrackedDetectedObject_)
+    @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', 'ymax', 'object_id', 'score', 'category', 'index']

dodgeballtracker/typing.py

@@ -4,40 +4,10 @@ import numpy as np
 import pandera as pa
 from pandera.dtypes import Float32, Float64
 
-XCoord = NewType('XCoord', float)
-YCoord = NewType('YCoord', float)
-Score = NewType('Confidence', float)
-Category = NewType('Category', float)
-
-
-@dataclass
-class Point:
-    x: XCoord
-    y: YCoord
-
-
 VideoFrame = NewType('Frame', float)
 VideoTimestamp = float
 
 
-@dataclass
-class DetectedObject:
-    xmin: XCoord
-    ymin: YCoord
-    xmax: XCoord
-    ymax: YCoord
-    score: Score
-    category: Category
-
-    def as_int_box(self) -> tuple[tuple[int, int], tuple[int, int]]:
-        return (int(self.xmin), int(self.ymin)), (
-            int(self.xmax), int(self.ymax))
-
-    def as_nparray(self):
-        return np.array([self.xmin, self.ymax, self.xmax, self.ymax, self.score,
-                         self.category])
-
-
 def tensor_of(data_type: pa.DataType, cols: list[str]) -> pa.DataFrameSchema:
     return pa.DataFrameSchema({col: pa.Column(data_type) for col in cols})
 
@@ -45,8 +15,8 @@ def tensor_of(data_type: pa.DataType, cols: list[str]) -> pa.DataFrameSchema:
 DetectedObject_ = tensor_of(Float32, ["xmin", "ymin", "xmax", "ymax", "score",
                                       "category"])
 
-TrackedDetectedObject_ = tensor_of(Float64,
-                                   list(DetectedObject_.columns) + ["object_id", "index"])
+TrackedDetectedObject = tensor_of(Float64,
+                                  list(DetectedObject_.columns) + ["object_id", "index"])
 # ["xmin", "ymin", "xmax", "ymax", "object_id", "score", "category", "index"])
 
 

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, DetectedObject, Color, DetectedObject_
+from dodgeballtracker.typing import VideoFrame, Color, DetectedObject_
 import matplotlib
 import cv2