From 2d0d09e6b02f0a92c2f82402b19813eef8bf8fea Mon Sep 17 00:00:00 2001 From: Mehul Goel Date: Sat, 4 Jan 2025 17:03:48 -0800 Subject: [PATCH] added benchmarking --- vision/ultralytics-benchmark.py | 41 +++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 vision/ultralytics-benchmark.py diff --git a/vision/ultralytics-benchmark.py b/vision/ultralytics-benchmark.py new file mode 100644 index 0000000..b83c3c2 --- /dev/null +++ b/vision/ultralytics-benchmark.py @@ -0,0 +1,41 @@ +import cv2 + +from ultralytics import YOLO + +# Load the YOLO11 model +model = YOLO("best.pt") + +# Open the video file +video_path = "data\9-21-whobaat1.avi" +cap = cv2.VideoCapture(video_path) + +subsample = 5 + +i = 0 +# Loop through the video frames +while cap.isOpened(): + # Read a frame from the video + success, frame = cap.read() + i += 1 + if (i % subsample != 0): + continue + if success: + # Run YOLO11 tracking on the frame, persisting tracks between frames + results = model.track(frame, persist=True) + + # Visualize the results on the frame + annotated_frame = results[0].plot() + + # Display the annotated frame + cv2.imshow("YOLO11 Tracking", annotated_frame) + + # Break the loop if 'q' is pressed + if cv2.waitKey(1) & 0xFF == ord("q"): + break + else: + # Break the loop if the end of the video is reached + break + +# Release the video capture object and close the display window +cap.release() +cv2.destroyAllWindows() \ No newline at end of file