import torch
import gradio as gr
from PIL import Image

# Load pre-trained YOLOv5 model from ultralytics
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

def detect_objects(image):
    # Perform inference
    results = model(image)
    results.render()  # updates results.imgs with boxes and labels
    detected_image = Image.fromarray(results.imgs[0])
    return detected_image

# Define Gradio interface
iface = gr.Interface(
    fn=detect_objects,
    inputs=gr.inputs.Image(type="pil"),
    outputs=gr.outputs.Image(type="pil"),
    title="Object Detection with YOLOv5",
    description="Upload an image and get the detected objects with their names."
)

# Launch the interface
iface.launch()