import gradio as gr
from ultralytics import YOLO
from PIL import Image

# Load YOLOv11 model (make sure yolo11n.pt is in the same folder)
model = YOLO("yolo11n.pt")

# Define prediction function
def detect_objects(image):
    results = model(image)
    # results[0].plot() returns a numpy image with detections drawn
    result_image = Image.fromarray(results[0].plot())
    return result_image

# Create Gradio interface
app = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="pil", label="Upload an image"),
    outputs=gr.Image(type="pil", label="Detections"),
    title="YOLOv11 Object Detection",
    description="Upload an image to detect objects using YOLOv11 Nano (yolo11n.pt)."
)

# Launch app
app.launch()