Commiting with problems in tracking.

dodgeballtracker/logic.py

@@ -1,16 +1,14 @@
 import json
-from collections import Counter
 from pathlib import Path
 import logging
 import torch
 from ultralytics import YOLO
 import cv2
 from dodgeballtracker.helper import iter_frames
-from dodgeballtracker.video_writer import create_video_writer, VideoWriter
-import matplotlib
+from dodgeballtracker.video_writer import VideoWriter, DetectedObjectBox
 import numpy as np
 from boxmot import DeepOcSort as DeepOCSORT
-from .typing import VideoFrame
+from .typing import VideoFrame, DetectedObject
 
 logger = logging.getLogger(__name__)
 
@@ -28,21 +26,6 @@ def get_best_available_torch_device() -> torch.device:
     return torch.device('cpu')
 
 
-class DetectedObjectBox:
-    def __init__(self, min_threshold: float, color: Color):
-        self.min_threshold = min_threshold
-        self.color = color
-
-    def __lt__(self, other):
-        return self.min_threshold < other.min_threshold
-
-    @classmethod
-    def from_viridis(cls, min_threshold):
-        cm = matplotlib.colormaps['viridis']
-        color: Color = tuple((int(255 * cm(min_threshold)[i]) for i in range(3)))
-        return cls(min_threshold, color)
-
-
 def get_court_perspective_transform():
     """
     """
@@ -117,13 +100,17 @@ class DodgeballTracker:
         for detected_box_object in self.detected_box_objects:
 
             if detected_box_object.min_threshold < confidence:
-                xmin, ymin, xmax, ymax = self.extract_box(data)
-                cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2)
+                # xmin, ymin, xmax, ymax = self.extract_box(data)
+                # cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2)
                 return True
         return False
 
     def run_model_on_frame(self, model, frame: VideoFrame):
-        return model(frame)[0]
+        model_outputs = model(frame)[0]
+        for data in model_outputs.boxes.data.tolist():
+            detected_obj = DetectedObject(*data)
+            if detected_obj.confidence > 0.5:
+                yield detected_obj
 
     def persist_points(self):
         with open(self.output_json_path, "wt") as file_:
@@ -142,30 +129,28 @@ 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]
+
     def main(self):
         # define some constants
         video_cap = cv2.VideoCapture(str(self.input_path))
         tracker = self._get_tracker()
         model = self._get_detection_model()
-        counter = Counter()
 
         for i, (frame, ts) in enumerate(iter_frames(video_cap)):
             # run the YOLO model on the frame
-            res = []
-            detections = self.run_model_on_frame(model, frame)
-            for data in detections.boxes.data.tolist():
-                is_good_box = self.plot_box_to_frame(data, frame)
-                if is_good_box:
-                    res.append(data)
-
-            data_ = tracker.update(np.array(res), frame)
-            self.print_tracker_data(data_, i, ts)
+            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([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)
 
             if self.video_writer.wait():
                 break
 
-        print(f"counter: {counter}")
         video_cap.release()
         self.persist_points()