going in the right direction.

This commit is contained in:
didier amyot committed 2024-01-21 08:51:30 -05:00
1 parent d716ba80f7
commit f23f17d428
1 file changed
+18 -2
+18 -2
View File
@@ -1,8 +1,9 @@
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
from numpy import arange
import numpy as np
Color = tuple[int, int, int]
@@ -45,6 +46,20 @@ 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)
@@ -54,6 +69,7 @@ def main(input_filename, output_filename, detected_box_objects):
# 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():
@@ -75,7 +91,7 @@ def main(input_filename, output_filename, detected_box_objects):
if __name__ == "__main__":
detected_box_objects = [
DetectedObjectBox.from_viridis(threshold) for threshold in arange(0.5, 1, 0.05)
DetectedObjectBox.from_viridis(threshold) for threshold in np.arange(0.5, 1, 0.05)
]
main("micro.mp4", "output.mp4", detected_box_objects)