made project vaguely in the shape of a python project

dodgeballtracker/logic.py

@@ -2,7 +2,7 @@ from pathlib import Path
 
 from ultralytics import YOLO
 import cv2
-from helper import create_video_writer, iter_frames, time_tracker
+from dodgeballtracker.helper import create_video_writer, iter_frames, time_tracker
 from deep_sort_realtime.deepsort_tracker import DeepSort
 import matplotlib
 import numpy as np

main.py

@@ -1,19 +0,0 @@
-import cv2 as cv
-
-minute = 60
-video_start = 9 * minute + 30
-video_end = 12 * minute + 56
-
-cap = cv.VideoCapture('out.mp4')
-while cap.isOpened():
-    ret, frame = cap.read()
-    # if frame is read correctly ret is True
-    if not ret:
-        print("Can't receive frame (stream end?). Exiting ...")
-        break
-    gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
-    cv.imshow('frame', gray)
-    if cv.waitKey(1) == ord('q'):
-        break
-cap.release()
-cv.destroyAllWindows()

object_tracking.py

@@ -1,97 +0,0 @@
-from ultralytics import YOLO
-import cv2
-from helper import create_video_writer, iter_frames, time_tracker
-from deep_sort_realtime.deepsort_tracker import DeepSort
-import matplotlib
-import numpy as np
-
-Color = tuple[int, int, int]
-
-
-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 write_frame(writer, frame):
-    cv2.imshow("Frame", frame)
-    writer.write(frame)
-
-
-def extract_box(data):
-    return int(data[0]), int(data[1]), int(data[2]), int(data[3])
-
-
-def plot_boxe_to_frame(data, frame, detected_box_objects):
-    confidence = float(data[4])
-    for detected_box_object in detected_box_objects:
-
-        if detected_box_object.min_threshold < confidence:
-            xmin, ymin, xmax, ymax = extract_box(data)
-            cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2)
-            return
-
-
-def run_model_on_frame(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 main(input_filename, output_filename, detected_box_objects):
-    # define some constants
-    detected_box_objects = sorted(detected_box_objects, reverse=True)
-
-    video_cap = cv2.VideoCapture(input_filename)
-    writer = create_video_writer(video_cap, output_filename)
-
-    # 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)
-            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()
-
-
-if __name__ == "__main__":
-    detected_box_objects = [
-        DetectedObjectBox.from_viridis(threshold) for threshold in np.arange(0.5, 1, 0.05)
-    ]
-
-    main("micro.mp4", "output.mp4", detected_box_objects)