diff --git a/.gitignore b/.gitignore index 0a764a4..d430e32 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,4 @@ env +*.mp4 +.idea +.DS_Store diff --git a/helper.py b/helper.py new file mode 100644 index 0000000..64ebd16 --- /dev/null +++ b/helper.py @@ -0,0 +1,34 @@ +import contextlib +import time +import cv2 + + +def create_video_writer(video_cap, output_filename): + + # grab the width, height, and fps of the frames in the video stream. + frame_width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + frame_height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + fps = int(video_cap.get(cv2.CAP_PROP_FPS)) + + # initialize the FourCC and a video writer object + fourcc = cv2.VideoWriter_fourcc(*'MP4V') + writer = cv2.VideoWriter(output_filename, fourcc, fps, + (frame_width, frame_height)) + + return writer + + +def iter_frames(video_cap): + while 1: + ret, frame = video_cap.read() + if not ret: + return + yield frame + + +@contextlib.contextmanager +def time_tracker(): + start = time.monotonic() + yield + end = time.monotonic() + print(f"Time to process 1 frame: {(end-start) * 1000:.0f} milliseconds") \ No newline at end of file diff --git a/main.py b/main.py index f24420b..3963ffb 100644 --- a/main.py +++ b/main.py @@ -1,5 +1,19 @@ import cv2 as cv -img = cv.imread("/Users/didier/ahuntsic/cours-420-301-ah/cours/images/json/001.jpg") -cv.imshow("Display window", img) -k = cv.waitKey(0) # Wait for a keystroke in the window +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 new file mode 100644 index 0000000..40f86ca --- /dev/null +++ b/object_tracking.py @@ -0,0 +1,81 @@ +from ultralytics import YOLO +import cv2 +from helper import create_video_writer, iter_frames, time_tracker +import matplotlib +from numpy import arange + +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 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") + + 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 arange(0.5, 1, 0.05) + ] + + main("micro.mp4", "output.mp4", detected_box_objects)