# import libraries
from ultralytics import YOLO
import cv2
import sys

def realtime(video):
    # Load the YOLOv8 model
    model = YOLO('yolov8n.pt')
    
    # Open the video file
    video_path = video
    cap = cv2.VideoCapture(video_path)
    cap.set(3, 720)
    cap.set(4, 1280)
    # 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
            results = model(frame, verbose=False)
    
            # Visualize the results on the frame
            annotated_frame = results[0].plot()
    
            # 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()
    
if __name__ == '__main__':
    realtime(sys.argv[1])