import cv2
import torch
from ultralytics import YOLO

# Load YOLOv5 model with GPU support if available
device = "cuda" if torch.cuda.is_available() else "cpu"
model = YOLO("yolov5s.pt").to(device)  # You can use yolov5m.pt for better accuracy

# Initialize video capture (Webcam)
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FPS, 30)  # Ensure 30+ FPS
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)  # Set width
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)  # Set height

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # Perform object detection
    results = model(frame)

    # Plot the results (draw bounding boxes)
    result_img = results[0].plot()

    # Display the output in a window
    cv2.imshow("YOLOv5 Live Object Detection", result_img)

    # Break loop with 'q' key
    if cv2.waitKey(1) & 0xFF == ord("q"):
        break

# Release resources
cap.release()
cv2.destroyAllWindows()