app.py

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()