import gradio as gr
import torch
import cv2
import numpy as np

# Load YOLOv5 model
def load_model():
    return torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

model = load_model()

# YOLOv5 object detection function
def yolo_object_detection(frame):
    # Convert Gradio's RGB frame to BGR for OpenCV processing
    img_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

    # Perform YOLOv5 inference
    results = model(img_bgr)
    detections = results.pandas().xyxy[0]

    # Draw bounding boxes and labels
    for _, row in detections.iterrows():
        x1, y1, x2, y2 = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax'])
        label = f"{row['name']} {row['confidence']:.2f}"
        cv2.rectangle(img_bgr, (x1, y1), (x2, y2), (255, 0, 0), 2)
        cv2.putText(img_bgr, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (36, 255, 12), 2)

    # Convert back to RGB for Gradio display
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    return img_rgb

# Gradio interface for real-time object detection
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            # Video input from webcam
            input_img = gr.Image(sources=["webcam"], type="numpy", label="Webcam Input")
        with gr.Column():
            # Output with detected objects
            output_img = gr.Image(streaming=True, label="YOLOv5 Object Detection Output")
        
        # Stream YOLOv5 object detection
        input_img.stream(yolo_object_detection, inputs=[input_img], outputs=[output_img],
                         time_limit=30, stream_every=0.1, concurrency_limit=30)

# Launch the app
demo.launch()