File size: 2,347 Bytes
3100b46
caff61e
bccf53b
dc80d48
3100b46
 
93fb6ed
2420aaa
b31672b
3100b46
 
a186d85
3c2cbbf
 
 
 
93fb6ed
8513c99
3c2cbbf
b31672b
3c2cbbf
 
b31672b
3100b46
 
93fb6ed
3100b46
b31672b
3100b46
 
 
3c2cbbf
b31672b
3c2cbbf
 
3100b46
 
3c2cbbf
2420aaa
3c2cbbf
3100b46
2420aaa
3100b46
 
 
3c2cbbf
3100b46
 
93fb6ed
3100b46
93fb6ed
3100b46
 
3c2cbbf
 
93fb6ed
 
b31672b
3c2cbbf
b31672b
3c2cbbf
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
import cv2
import torch
import numpy as np
import gradio as gr
from ultralytics import YOLO
import threading
import time

# Load YOLOv5 model
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = YOLO("yolov5s.pt").to(device)

# Open webcam
cap = cv2.VideoCapture(0)
frame = np.zeros((480, 640, 3), dtype=np.uint8)  # Default blank frame
lock = threading.Lock()

def detect_objects(image):
    """Detect objects in an uploaded image with YOLO."""
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # Convert to BGR for OpenCV
    results = model.predict(image, conf=0.4)  # Set confidence threshold
    image = results[0].plot()  # Plot detections directly on image
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert back to RGB for Gradio

def process_webcam():
    """Continuously capture and process frames from the webcam."""
    global frame
    while cap.isOpened():
        ret, img = cap.read()
        if not ret:
            continue

        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # Convert to RGB
        results = model.predict(img, conf=0.4)  # Explicitly call predict
        img = results[0].plot()  # Directly draw detections on the frame
        
        with lock:
            frame = img  # Update frame with detection overlay

# Start the webcam processing thread
threading.Thread(target=process_webcam, daemon=True).start()

def get_webcam_frame():
    """Returns the latest processed webcam frame."""
    with lock:
        return frame

# 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", type="numpy")

            def update_webcam():
                while True:
                    with lock:
                        img = frame
                    webcam_output.update(img)
                    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()