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