Spaces:
Sleeping
Sleeping
File size: 1,537 Bytes
908e980 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
import os
import cv2
from ultralytics import YOLO
PROJECT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
yolo_model = os.path.join(PROJECT_DIR, "output", "yolov8n_custom", "weights", "best.pt")
# Load the YOLOv8 model
model = YOLO(yolo_model)
# Open the video file
video_path = os.path.join(PROJECT_DIR, "data_for_inference", "video", "vid_39_1284-2_11.mp4")
cap = cv2.VideoCapture(video_path)
display_size = (800, 600)
# Loop through the video frames
while cap.isOpened():
# Read a frame from the video
success, frame = cap.read()
if success:
# Run YOLOv8 inference on the frame
# save plotted images
plotted_result_path = os.path.join(PROJECT_DIR, "scripts", "inference")
results = model.predict(source=frame,
save=True,
conf=0.2,
project=plotted_result_path)
# Visualize the results on the frame
annotated_frame = results[0].plot()
# resize the frame to fit the display window
annotated_frame = cv2.resize(annotated_frame, display_size)
# Display the annotated frame
cv2.imshow("YOLOv8 Inference", 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()
|