Spaces:
Sleeping
Sleeping
| 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() | |