navjotk's picture
Update app.py
90a129b verified
raw
history blame
697 Bytes
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()