app.py

import gradio as gr
from transformers import pipeline
from PIL import Image, ImageDraw, ImageFont
import tempfile

# Load the YOLOS object detection model

detector = pipeline("object-detection", model="hustvl/yolos-tiny")

# 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):
    results = detector(image)
    image = image.convert("RGB")
    draw = ImageDraw.Draw(image)

    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.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))

    # Save image for download
    temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    image.save(temp_file.name)

    # ✅ Return the (image, annotations) tuple and the path to the saved image
    return (image, annotations), temp_file.name

# 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"),
        gr.File(label="Download Processed Image"),
    ],
    title="YOLOS Object Detection",
    description="Upload an image to detect objects using the YOLOS-small model. Adjust the confidence threshold using the slider.",
)

demo.launch()