Spaces:
Running
Running
File size: 2,234 Bytes
a09e5e0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import gradio as gr
from transformers import pipeline
from PIL import Image, ImageDraw, ImageFont
# Load the YOLOS object detection model
detector = pipeline("object-detection", model="hustvl/yolos-small")
# Define some colors to differentiate classes
COLORS = ["red", "blue", "green", "orange", "purple", "yellow", "cyan", "magenta"]
# Helper function to assign color per label
def get_color_for_label(label):
return COLORS[hash(label) % len(COLORS)]
# Main function: detect, draw, and return outputs
def detect_and_draw(image, threshold):
try:
# Perform object detection
results = detector(image)
image = image.convert("RGB")
draw = ImageDraw.Draw(image)
# Try to load a font for annotations, else use default
try:
font = ImageFont.truetype("arial.ttf", 16)
except:
font = ImageFont.load_default()
annotations = []
for obj in results:
score = obj["score"]
if score < threshold:
continue
label = f"{obj['label']} ({score:.2f})"
box = obj["box"]
color = get_color_for_label(obj["label"])
# Draw the bounding box and label
draw.rectangle(
[(box["xmin"], box["ymin"]), (box["xmax"], box["ymax"])],
outline=color,
width=3,
)
draw.text((box["xmin"] + 5, box["ymin"] + 5), label, fill=color, font=font)
box_coords = (box["xmin"], box["ymin"], box["xmax"], box["ymax"])
annotations.append((box_coords, label))
return image, annotations
except Exception as e:
return f"Error during detection: {e}", None
# Gradio UI setup
demo = gr.Interface(
fn=detect_and_draw,
inputs=[
gr.Image(type="pil", label="Upload Image"),
gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.05, label="Confidence Threshold"),
],
outputs=[
gr.AnnotatedImage(label="Detected Image"),
],
title="YOLOS Object Detection",
description="Upload an image to detect objects using the YOLOS-small model. Adjust the confidence threshold using the slider.",
live=True
)
demo.launch()
|