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
|
||||
*.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
|
||||
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()
|
||||
@@ -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