app.py

import os
import gradio as gr
from PIL import Image
from ultralytics import YOLO
from spaces import GPU  # Enables GPU-specific execution in Hugging Face Spaces

# Prevent config write errors by setting Ultralytics settings path to /tmp
os.environ["YOLO_CONFIG_DIR"] = "/tmp"

# Load YOLO11 model (you can also use yolo11s.pt, yolo11m.pt etc.)
model = YOLO("yolo11n.pt")

@GPU  # Ensures GPU is allocated for this function
def detect_objects(image: Image.Image) -> Image.Image:
    results = model(image)
    return results[0].plot()  # Returns image with bounding boxes drawn

# Create Gradio interface
demo = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="pil"),
    outputs="image",
    title="YOLO11 Object Detection",
    description="Upload an image to detect objects using Ultralytics YOLO11"
)

# Launch app
demo.launch()