.gitignore

@@ -160,3 +160,4 @@ env
 *.png
 *.ipynb
 .direnv
+result

dodgeballtracker/logic.py

@@ -107,9 +107,12 @@ class DodgeballTracker:
 
     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.confidence > 0.5:
+            if detected_obj.score > 0.5:
                 yield detected_obj
 
     def persist_points(self):
@@ -130,7 +133,7 @@ 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.confidence > 0.5]
+        important_detected_objs = [do for do in detected_objs if do.score > 0.5]
 
     def main(self):
         # define some constants
@@ -141,7 +144,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))
-            self.video_writer.plot_boxes_to_frame(detected_objs, frame)
+            data_ = tracker.update(
+                np.array(detected_objs), frame)
+            if i == 100:
+                print(data_)
+            # 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")

dodgeballtracker/typing.py

@@ -4,7 +4,7 @@ import numpy as np
 
 XCoord = NewType('XCoord', float)
 YCoord = NewType('YCoord', float)
-Confidence = NewType('Confidence', float)
+Score = NewType('Confidence', float)
 Category = NewType('Category', float)
 
 
@@ -23,14 +23,14 @@ class DetectedObject:
     ymin: YCoord
     xmax: XCoord
     ymax: YCoord
-    confidence: Confidence
+    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.confidence, self.category])
+        return np.array([self.xmin, self.ymax, self.xmax, self.ymax, self.score, self.category])
 
 
 

dodgeballtracker/video_writer.py

@@ -81,7 +81,7 @@ class MP4VVideoWriter(VideoWriter):
         self._video_writer.write(frame)
 
     def plot_box_to_frame(self, data: DetectedObject, frame: VideoFrame) -> None:
-        confidence = float(data.confidence)
+        confidence = float(data.score)
         for threshold in self._thresholds:
             if threshold.min_threshold < confidence:
                 # xmin, ymin, xmax, ymax = self._extract_box(data)

flake.nix

@@ -22,8 +22,42 @@
           system = system;
           config.allowUnfree = true;
         };
+
         pythonPackages = pkgs.python312Packages
           // opencv-python-related.packages.${system}.python312Packages;
+
+        pandera = pkgs.python3Packages.buildPythonPackage  {
+          version = "0.5.11";
+          pname = "lapx";
+          pyproject = true;
+          src = pkgs.python3Packages.fetchPypi {
+            pname = "pandera";
+            version = "0.20.4";
+            hash = "sha256-zPYXgpPvnUOT3Ed25HR349LdUcgA7N/t7Gf/9Q9K08I=";
+
+          };
+
+          propagatedBuildInputs = with pythonPackages; [ multimethod
+            numpy
+            pandas
+            pydantic
+            typeguard
+            typing-inspect
+            wrapt
+           ];
+
+          nativeBuildInputs = with pythonPackages; [
+            setuptools
+          ];
+
+          meta = with pkgs.lib; {
+            description = "The Open-source Framework for Precision Data Testing";
+            homepage = "https://github.com/unionai-oss/pandera";
+            license = pkgs.lib.licenses.mit;
+            maintainers = with maintainers; [ ];
+          };
+        };
+
         virtual-env = pkgs.python312.withPackages (ps: [
           ps.matplotlib
           ps.torchWithCuda
@@ -32,10 +66,11 @@
           pythonPackages.boxmot
           pythonPackages.ultralytics
           pythonPackages.click
+          pandera
         ]);
 
       in {
-        packages = { env = virtual-env; };
+        packages = { env = virtual-env; pandera = pandera; };
         devShells.default = pkgs.mkShell {
           buildInputs = [ virtual-env pkgs.yt-dlp pkgs.ffmpeg_7 ];
         };