going in the very right direction!

dodgeballtracker/logic.py

@@ -1,3 +1,5 @@
+import json
+import sys
 from pathlib import Path
 
 from ultralytics import YOLO
@@ -34,7 +36,10 @@ class DodgeballTracker:
                  detected_box_objects: list[DetectedObjectBox]):
         self.input_path = input_file
         self.output_path = output_file
+        self.output_json_path = output_file.with_suffix(".json")
         self.detected_box_objects = sorted(detected_box_objects, reverse=True)
+        self.transformation_matrix = self.get_court_perspective_transform()
+        self.points = []
 
     def write_frame(self, writer, frame) -> None:
         cv2.imshow("Frame", frame)
@@ -43,18 +48,6 @@ class DodgeballTracker:
     def extract_box(self, data) -> tuple[int, int, int, int]:
         return int(data[0]), int(data[1]), int(data[2]), int(data[3])
 
-    def plot_boxe_to_frame(self, data, frame):
-        confidence = float(data[4])
-        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)
-                return
-
-    def run_model_on_frame(self, model, frame):
-        return model(frame)[0]
-
     def get_court_perspective_transform(self):
         """
         """
@@ -79,6 +72,30 @@ class DodgeballTracker:
                                        court_bottom_right]) * pixel_per_meter
         return cv2.getPerspectiveTransform(empirical_source, empirical_target)
 
+    def print_box(self, data, frame_idx, frame_ts):
+        xmin, ymin, xmax, ymax = self.extract_box(data)
+        v = np.array([
+            (xmax + xmin) / 2,
+            ymax,  # (ymax + ymin) / 2,
+            1
+        ])
+        ret = self.transformation_matrix @ v
+        x, y = ret[0] / ret[2], ret[1] / ret[2]
+        self.points.append({"x": x, "y": y, "frameIndex": frame_idx, "frameTimestamp": frame_ts})
+
+    def plot_boxe_to_frame(self, data, frame, frame_idx, frame_ts):
+        confidence = float(data[4])
+        for detected_box_object in self.detected_box_objects:
+
+            if detected_box_object.min_threshold < confidence:
+                xmin, ymin, xmax, ymax = self.extract_box(data)
+                self.print_box(data, frame_idx, frame_ts)
+                cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2)
+                return
+
+    def run_model_on_frame(self, model, frame):
+        return model(frame)[0]
+
     def crazy_experiment(self, frame):
         cv2.putText(frame,
                     "Hello world",
@@ -92,6 +109,10 @@ class DodgeballTracker:
 
         cv2.polylines(frame, pts, True, (0, 255, 255), 8)
 
+    def persist_points(self):
+        with open(self.output_json_path, "wt") as file_:
+            json.dump({"points": self.points}, file_, indent=4)
+
     def main(self):
         # define some constants
 
@@ -101,22 +122,22 @@ class DodgeballTracker:
         # load the pre-trained YOLOv8n model
         model = YOLO("yolov8n.pt")
 
-        for frame in iter_frames(video_cap):
-            with time_tracker():
-                # run the YOLO model on the frame
-                detections = self.run_model_on_frame(model, frame)
-                for data in detections.boxes.data.tolist():
-                    self.plot_boxe_to_frame(data, frame)
+        for i, (frame, ts) in enumerate(iter_frames(video_cap)):
+            # run the YOLO model on the frame
+            detections = self.run_model_on_frame(model, frame)
+            for data in detections.boxes.data.tolist():
+                self.plot_boxe_to_frame(data, frame, i, ts)
 
-                # show the frame to our screen
-                self.write_frame(writer, frame)
+            # show the frame to our screen
+            self.write_frame(writer, frame)
 
-                if cv2.waitKey(1) == ord("q"):
-                    break
+            if cv2.waitKey(1) == ord("q"):
+                break
 
         video_cap.release()
         writer.release()
         cv2.destroyAllWindows()
+        self.persist_points()