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,44 +26,36 @@ class DetectedObjectBox:
         return cls(min_threshold, color)
 
 
-def write_frame(writer, frame):
+class DodgeballTracker:
+
+    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 write_frame(self, writer, frame) -> None:
         cv2.imshow("Frame", frame)
         writer.write(frame)
 
-
-def extract_box(data):
+    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])
-    for detected_box_object in detected_box_objects:
+        for detected_box_object in self.detected_box_objects:
 
             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)
                 return
 
-
-def run_model_on_frame(model, frame):
+    def run_model_on_frame(self, model, frame):
         return model(frame)[0]
 
-
-def crazy_experiment(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 get_court_perspective_transform(self):
         """
         """
         empirical_top_left = np.float32([416, 175])
@@ -87,27 +79,37 @@ def get_court_perspective_transform():
                                        court_bottom_right]) * pixel_per_meter
         return cv2.getPerspectiveTransform(empirical_source, empirical_target)
 
+    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))
 
-def main(input_filename: Path, output_filename: Path, detected_box_objects):
+        cv2.polylines(frame, pts, True, (0, 255, 255), 8)
+
+    def main(self):
         # define some constants
-    detected_box_objects = sorted(detected_box_objects, reverse=True)
 
-    video_cap = cv2.VideoCapture(str(input_filename))
-    writer = create_video_writer(video_cap, str(output_filename))
+        video_cap = cv2.VideoCapture(str(self.input_path))
+        writer = create_video_writer(video_cap, str(self.output_path))
 
         # load the pre-trained YOLOv8n model
         model = YOLO("yolov8n.pt")
-    tracker = DeepSort(max_age=50)
 
         for frame in iter_frames(video_cap):
             with time_tracker():
                 # run the YOLO model on the frame
-            detections = run_model_on_frame(model, frame)
+                detections = self.run_model_on_frame(model, frame)
                 for data in detections.boxes.data.tolist():
-                plot_boxe_to_frame(data, frame, detected_box_objects)
+                    self.plot_boxe_to_frame(data, frame)
 
                 # show the frame to our screen
-            write_frame(writer, frame)
+                self.write_frame(writer, frame)
 
                 if cv2.waitKey(1) == ord("q"):
                     break