diff --git a/object_tracking.py b/object_tracking.py index 40f86ca..8f90c69 100644 --- a/object_tracking.py +++ b/object_tracking.py @@ -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)