Spaces:
Sleeping
Sleeping
File size: 2,925 Bytes
3100b46 caff61e bccf53b dc80d48 3100b46 93fb6ed 2420aaa 3100b46 a186d85 93fb6ed 8513c99 93fb6ed 3100b46 df054c6 3100b46 93fb6ed 2420aaa 3100b46 93fb6ed 3100b46 2420aaa 3100b46 2420aaa 3100b46 93fb6ed 3100b46 2420aaa 3100b46 2420aaa 3100b46 93fb6ed 3100b46 93fb6ed 3100b46 93fb6ed 3100b46 93fb6ed 3100b46 8513c99 3100b46 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import cv2
import torch
import numpy as np
import gradio as gr
from ultralytics import YOLO
import threading
import time
# Load YOLOv5 model (optimized for CUDA if available)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = YOLO("yolov5s.pt").to(device)
# Generate unique colors for each class
num_classes = len(model.names)
colors = np.random.randint(0, 255, size=(num_classes, 3), dtype=np.uint8)
def detect_objects(image):
"""Detect objects in an uploaded image with different bounding box colors."""
results = model(image)
detections = results[0].boxes.data.cpu().numpy() # Get detections
for box in detections:
x1, y1, x2, y2, conf, cls = map(int, box[:6])
label = f"{model.names[cls]} {conf:.2f}"
color = tuple(map(int, colors[cls])) # Assign unique color based on class
cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
return image
# Real-time webcam processing
cap = cv2.VideoCapture(0) # Capture from webcam
frame = None
lock = threading.Lock()
def process_webcam():
"""Continuously capture and process frames from the webcam."""
global frame
while True:
ret, img = cap.read()
if not ret:
continue
results = model(img)
detections = results[0].boxes.data.cpu().numpy()
for box in detections:
x1, y1, x2, y2, conf, cls = map(int, box[:6])
label = f"{model.names[cls]} {conf:.2f}"
color = tuple(map(int, colors[cls])) # Assign unique color
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
with lock:
frame = img
# Start the webcam thread
threading.Thread(target=process_webcam, daemon=True).start()
def get_webcam_frame():
"""Returns the latest processed webcam frame."""
with lock:
return frame if frame is not None else np.zeros((480, 640, 3), dtype=np.uint8)
# Gradio UI
with gr.Blocks() as demo:
gr.Markdown("# YOLOv5 Real-Time Object Detection")
with gr.Tabs():
with gr.Tab("Real-Time Webcam"):
webcam_output = gr.Image(label="Live Webcam Feed")
def update_webcam():
while True:
webcam_output.update(get_webcam_frame())
time.sleep(1/30) # ~30 FPS
threading.Thread(target=update_webcam, daemon=True).start()
with gr.Tab("Upload Image"):
image_input = gr.Image(type="numpy", label="Upload Image")
image_output = gr.Image(label="Detected Objects")
image_button = gr.Button("Detect Objects")
image_button.click(detect_objects, inputs=image_input, outputs=image_output)
demo.launch()
|