From e26489bf9c251aa48377195715f7cec7473eb0bc Mon Sep 17 00:00:00 2001 From: didier amyot Date: Sun, 21 Jan 2024 09:04:43 -0500 Subject: [PATCH] made project vaguely in the shape of a python project --- helper.py => dodgeballtracker/helper.py | 0 dodgeballtracker/logic.py | 2 +- main.py | 19 ----- object_tracking.py | 97 ------------------------- 4 files changed, 1 insertion(+), 117 deletions(-) rename helper.py => dodgeballtracker/helper.py (100%) delete mode 100644 main.py delete mode 100644 object_tracking.py diff --git a/helper.py b/dodgeballtracker/helper.py similarity index 100% rename from helper.py rename to dodgeballtracker/helper.py diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py index 59cbe35..1b9362c 100644 --- a/dodgeballtracker/logic.py +++ b/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 diff --git a/main.py b/main.py deleted file mode 100644 index 3963ffb..0000000 --- a/main.py +++ /dev/null @@ -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() diff --git a/object_tracking.py b/object_tracking.py deleted file mode 100644 index 8f90c69..0000000 --- a/object_tracking.py +++ /dev/null @@ -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)