import gradio as gr
import torch
import cv2
import numpy as np
from PIL import Image

# Load the YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

# Function to perform object detection
def detect_objects(image):
    # Convert the image to a numpy array
    img_array = np.array(image)

    # Perform object detection
    results = model(img_array)

    # Draw bounding boxes and labels
    for index, row in results.pandas().xyxy[0].iterrows():
        x1, y1, x2, y2 = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax'])
        label = row['name']
        cv2.rectangle(img_array, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(img_array, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)

    # Convert the numpy array back to an image
    output_image = Image.fromarray(img_array)
    return output_image

# Set up the Gradio interface
interface = gr.Interface(
    fn=detect_objects,
    inputs=gr.inputs.Image(type="pil", tool="editor"),
    outputs=gr.outputs.Image(type="pil"),
    live=True,
    title="Real-time Object Detection with YOLOv5"
)

# Launch the interface
interface.launch()