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

    # Return the annotated image and annotations (no download option)
    return image, annotations

# 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.",
)

demo.launch()