# 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()