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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).to(device)

if device.type == 'cuda':
    model.half()

def detect_objects(image):
    image_pil = Image.fromarray(image)
    
    with torch.no_grad():
        results = model(image_pil)
    
    rendered_images = results.render()
    
    return np.array(rendered_images[0]) if rendered_images else image

# Gradio interface
iface = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=gr.Image(type="numpy", label="Detected Objects"),
    title="Object Detection with YOLOv5",
    description="Use webcam or upload an image to detect objects.",
    allow_flagging="never",
    examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
)

iface.launch()