Going somewhere before going to the gym

dodgeballtracker/logic.py

@@ -6,6 +6,7 @@ from ultralytics import YOLO
 import cv2
 from dodgeballtracker.helper import iter_frames
 from dodgeballtracker.video_writer import VideoWriter, DetectedObjectBox
+from dodgeballtracker.data_dumper import DataDumper
 import numpy as np
 import pandas as pd
 import pandera as pa
@@ -60,52 +61,17 @@ class DodgeballTracker:
 
     def __init__(self,
                  input_file: Path,
-                 output_file: Path,
                  video_writer: VideoWriter,
+                 data_dumper: DataDumper,
                  detected_box_objects: list[DetectedObjectBox],
                  transformation_matrix):
         self.input_path = input_file
-        self.output_path = output_file
-        self.output_json_path = output_file.with_suffix(".json")
         self.video_writer = video_writer
+        self._data_dumper = data_dumper
         self.detected_box_objects = sorted(detected_box_objects, reverse=True)
         self.transformation_matrix = transformation_matrix
         self.points = []
 
-    def write_frame(self, writer, frame) -> None:
-        cv2.imshow("Frame", frame)
-        writer.write(frame)
-
-    def extract_box(self, data) -> tuple[int, int, int, int]:
-        return int(data[0]), int(data[1]), int(data[2]), int(data[3])
-
-    def print_box(self, data, frame_idx, frame_ts):
-        xmin, ymin, xmax, ymax = self.extract_box(data)
-        xavg = (xmax + xmin) / 2
-        v = np.array([xavg, ymax, 1])
-        ret = self.transformation_matrix @ v
-        x, y = ret[0] / ret[2], ret[1] / ret[2]
-        object_id = data[-1]
-        self.points.append({"x": x,
-                            "y": y,
-                            "frameIndex": frame_idx,
-                            "frameTimestamp": frame_ts,
-                            "objectId": int(object_id)
-                            })
-
-    def print_tracker_data(self, tracker_data, frame_idx, frame_ts):
-        for d in tracker_data:
-            self.print_box(d, frame_idx, frame_ts)
-
-    def plot_box_to_frame(self, data, frame: VideoFrame):
-        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 True
-        return False
 
     @pa.check_output(DetectedObject)
     def run_model_on_frame(self, model, frame: VideoFrame):
@@ -115,9 +81,6 @@ class DodgeballTracker:
         ret = pd.DataFrame(boxes, columns=columns)
         return ret
 
-    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) -> DeepOCSORT:
         return DeepOCSORT(
@@ -155,6 +118,7 @@ class DodgeballTracker:
             # run the YOLO model on the frame
             detected_objs = self.run_model_on_frame(model, frame)
             tracked_obj = self.update_tracker(tracker, detected_objs, frame)
+            self._data_dumper.write_boxes(tracked_obj, ts, i)
             self.video_writer.plot_boxes_to_frame(tracked_obj, frame)
             self.video_writer.write_frame(frame)
 
@@ -162,4 +126,3 @@ class DodgeballTracker:
                 break
 
         video_cap.release()
-        self.persist_points()