Spaces:
Sleeping
Sleeping
File size: 1,810 Bytes
08d2f61 |
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 |
# app.py
from transformers import pipeline
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
import random
# Load the YOLO-based object detection pipeline
detector = pipeline("object-detection", model="hustvl/yolos-tiny")
# Generate a random color for each label
label_colors = {}
def get_color(label):
if label not in label_colors:
label_colors[label] = (
random.randint(0, 255),
random.randint(0, 255),
random.randint(0, 255)
)
return label_colors[label]
# Detection function
def detect_objects(img):
results = detector(img)
draw = ImageDraw.Draw(img)
font = ImageFont.load_default()
for obj in results:
label = obj["label"]
score = obj["score"]
box = obj["box"]
color = get_color(label)
# Draw bounding box
draw.rectangle(
[box["xmin"], box["ymin"], box["xmax"], box["ymax"]],
outline=color,
width=3
)
# Prepare label with confidence
label_text = f"{label} ({score:.2f})"
text_bbox = draw.textbbox((box["xmin"], box["ymin"]), label_text, font=font)
text_background = [text_bbox[0], text_bbox[1], text_bbox[2], text_bbox[3]]
# Draw background for text
draw.rectangle(text_background, fill=color)
draw.text((text_bbox[0], text_bbox[1]), label_text, fill="black", font=font)
return img
# Gradio interface
interface = gr.Interface(
fn=detect_objects,
inputs=gr.Image(type="pil"),
outputs=gr.Image(type="pil"),
title="YOLO Object Detection with Color-coded Labels",
description="Upload an image. Detected objects are shown with bounding boxes and color-coded labels using YOLOS-Tiny."
)
if __name__ == "__main__":
interface.launch()
|