migrated to class

dodgeballtracker/__init__.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 
-from dodgeballtracker.logic import DetectedObjectBox, main
+from dodgeballtracker.logic import DetectedObjectBox, DodgeballTracker
 import numpy as np
 
 
@@ -12,8 +12,9 @@ if __name__ == "__main__":
     detected_box_objects = [
         DetectedObjectBox.from_viridis(threshold) for threshold in np.arange(0.5, 1, 0.05)
     ]
-
     input_video = input_dir / "micro.mp4"
     output_video = output_dir / "output.mp4"
 
-    main(input_video, output_video, detected_box_objects)
+    dodgeball_tracker = DodgeballTracker(input_video, output_video, detected_box_objects)
+
+    dodgeball_tracker.main()

dodgeballtracker/logic.py

@@ -26,95 +26,97 @@ class DetectedObjectBox:
         return cls(min_threshold, color)
 
 
-def write_frame(writer, frame):
+class DodgeballTracker:
-    cv2.imshow("Frame", frame)
-    writer.write(frame)
 
+    def __init__(self,
+                 input_file: Path,
+                 output_file: Path,
+                 detected_box_objects: list[DetectedObjectBox]):
+        self.input_path = input_file
+        self.output_path = output_file
+        self.detected_box_objects = sorted(detected_box_objects, reverse=True)
 
-def extract_box(data):
+    def write_frame(self, writer, frame) -> None:
-    return int(data[0]), int(data[1]), int(data[2]), int(data[3])
+        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 plot_boxe_to_frame(data, frame, detected_box_objects):
+    def plot_boxe_to_frame(self, data, frame):
-    confidence = float(data[4])
+        confidence = float(data[4])
-    for detected_box_object in detected_box_objects:
+        for detected_box_object in self.detected_box_objects:
 
-        if detected_box_object.min_threshold < confidence:
+            if detected_box_object.min_threshold < confidence:
-            xmin, ymin, xmax, ymax = extract_box(data)
+                xmin, ymin, xmax, ymax = self.extract_box(data)
-            cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2)
+                cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2)
-            return
+                return
 
+    def run_model_on_frame(self, model, frame):
+        return model(frame)[0]
 
-def run_model_on_frame(model, frame):
+    def get_court_perspective_transform(self):
-    return model(frame)[0]
+        """
+        """
+        empirical_top_left = np.float32([416, 175])
+        empirical_top_right = np.float32([905, 173])
+        empirical_bottom_left = np.float32([51, 702])
+        empirical_bottom_right = np.float32([1188, 695])
 
+        court_top_left = np.float32([0, 0])
+        court_top_right = np.float32([9, 0])
+        court_bottom_left = np.float32([0, 18])
+        court_bottom_right = np.float32([9, 18])
 
-def crazy_experiment(frame):
+        empirical_source = np.float32([empirical_top_left,
-    cv2.putText(frame,
+                                       empirical_top_right,
-                "Hello world",
+                                       empirical_bottom_left,
-                (50, 50),
+                                       empirical_bottom_right])
-                cv2.FONT_HERSHEY_SIMPLEX,
+        pixel_per_meter = 50
-                2,
+        empirical_target = np.float32([court_top_left,
-                (0, 255, 0), 8)
+                                       court_top_right,
-    cv2.line(frame, (0, 0), (511, 511), (255, 0, 0), 5)
+                                       court_bottom_left,
-    pts = np.array([[50, 50], [500, 50], [500, 500], [50, 500], [50, 50]], np.int32)
+                                       court_bottom_right]) * pixel_per_meter
-    pts = pts.reshape((-1, 1, 2))
+        return cv2.getPerspectiveTransform(empirical_source, empirical_target)
 
-    cv2.polylines(frame, pts, True, (0, 255, 255), 8)
+    def crazy_experiment(self, frame):
+        cv2.putText(frame,
+                    "Hello world",
+                    (50, 50),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    2,
+                    (0, 255, 0), 8)
+        cv2.line(frame, (0, 0), (511, 511), (255, 0, 0), 5)
+        pts = np.array([[50, 50], [500, 50], [500, 500], [50, 500], [50, 50]], np.int32)
+        pts = pts.reshape((-1, 1, 2))
 
+        cv2.polylines(frame, pts, True, (0, 255, 255), 8)
 
-def get_court_perspective_transform():
+    def main(self):
-    """
+        # define some constants
-    """
-    empirical_top_left = np.float32([416, 175])
-    empirical_top_right = np.float32([905, 173])
-    empirical_bottom_left = np.float32([51, 702])
-    empirical_bottom_right = np.float32([1188, 695])
 
-    court_top_left = np.float32([0, 0])
+        video_cap = cv2.VideoCapture(str(self.input_path))
-    court_top_right = np.float32([9, 0])
+        writer = create_video_writer(video_cap, str(self.output_path))
-    court_bottom_left = np.float32([0, 18])
-    court_bottom_right = np.float32([9, 18])
 
-    empirical_source = np.float32([empirical_top_left,
+        # load the pre-trained YOLOv8n model
-                                  empirical_top_right,
+        model = YOLO("yolov8n.pt")
-                                  empirical_bottom_left,
-                                  empirical_bottom_right])
-    pixel_per_meter = 50
-    empirical_target = np.float32([court_top_left,
-                                  court_top_right,
-                                  court_bottom_left,
-                                  court_bottom_right]) * pixel_per_meter
-    return cv2.getPerspectiveTransform(empirical_source, empirical_target)
 
+        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)
 
-def main(input_filename: Path, output_filename: Path, detected_box_objects):
+                # show the frame to our screen
-    # define some constants
+                self.write_frame(writer, frame)
-    detected_box_objects = sorted(detected_box_objects, reverse=True)
 
-    video_cap = cv2.VideoCapture(str(input_filename))
+                if cv2.waitKey(1) == ord("q"):
-    writer = create_video_writer(video_cap, str(output_filename))
+                    break
 
-    # load the pre-trained YOLOv8n model
+        video_cap.release()
-    model = YOLO("yolov8n.pt")
+        writer.release()
-    tracker = DeepSort(max_age=50)
+        cv2.destroyAllWindows()
-
-    for frame in iter_frames(video_cap):
-        with time_tracker():
-            # run the YOLO model on the frame
-            detections = run_model_on_frame(model, frame)
-            for data in detections.boxes.data.tolist():
-                plot_boxe_to_frame(data, frame, detected_box_objects)
-
-            # show the frame to our screen
-            write_frame(writer, frame)
-
-            if cv2.waitKey(1) == ord("q"):
-                break
-
-    video_cap.release()
-    writer.release()
-    cv2.destroyAllWindows()