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

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

# Load YOLOv5s model (smallest and fastest variant)
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).to(device)

# Enable half-precision for CUDA devices
if device.type == 'cuda':
    model.half()

def detect_objects(image):
    # Convert numpy array to PIL Image
    image_pil = Image.fromarray(image)
    
    # Perform inference without gradient calculation
    with torch.no_grad():
        results = model(image_pil)
    
    # Render detections using optimized YOLOv5 method
    rendered_images = results.render()
    
    # Return the first rendered image as numpy array
    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="High-Speed Object Detection with YOLOv5s",
    description="Optimized for speed using YOLOv5s and GPU acceleration.",
    allow_flagging="never",
    examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
)

iface.launch()