moved files

dodgeballtracker/logic.py

@@ -0,0 +1,92 @@
+from pathlib import Path
+
+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: Path, output_filename: Path, detected_box_objects):
+    # 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))
+
+    # 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()
+