import gradio as gr
from transformers import pipeline
from PIL import Image

# Load YOLOS pipeline
detector = pipeline("object-detection", model="hustvl/yolos-small")

def detect_objects(img):
    results = detector(img)
    boxes = []
    for obj in results:
        label = f"{obj['label']} ({obj['score']:.2f})"
        box = obj['box']
        boxes.append((box["xmin"], box["ymin"], box["xmax"], box["ymax"], label))
    return img, boxes

demo = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="pil"),
    outputs=gr.AnnotatedImage(),
    title="YOLOS Object Detection",
    description="Upload an image and detect objects using the YOLOS Transformer model.",
)

demo.launch()