going in the right direction.
This commit is contained in:
1 parent
b69b33ee59
commit
d716ba80f7
4 files changed
+135
-3
No files matched your search
@@ -1 +1,4 @@
|
|||||||
env
|
env
|
||||||
|
*.mp4
|
||||||
|
.idea
|
||||||
|
.DS_Store
|
||||||
@@ -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")
|
||||||
@@ -1,5 +1,19 @@
|
|||||||
import cv2 as cv
|
import cv2 as cv
|
||||||
img = cv.imread("/Users/didier/ahuntsic/cours-420-301-ah/cours/images/json/001.jpg")
|
|
||||||
|
|
||||||
cv.imshow("Display window", img)
|
minute = 60
|
||||||
k = cv.waitKey(0) # Wait for a keystroke in the window
|
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()
|
||||||
@@ -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)
|
||||||
Reference in new issue
Block a user