from ultralytics import YOLO
import cv2
import sys

# Check for command-line argument
if len(sys.argv) < 2:
    print("Usage: python detect_video.py <video_path>")
    exit()

video_path = sys.argv[1]

model = YOLO("best50.pt")
cap = cv2.VideoCapture(video_path)

if not cap.isOpened():
    print(f"Error: Could not open video file {video_path}.")
    exit()

print("Press 'q' to quit.")
while True:
    ret, frame = cap.read()
    if not ret:
        print("End of video or failed to grab frame.")
        break

    results = model(frame, conf=0.5, iou=0.4, imgsz=640, augment=False)
    annotated_frame = results[0].plot()

    cv2.imshow("Video Detection", annotated_frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()
print("Video closed.")