import gradio as gr  # ← 추가!
from transformers import pipeline
from PIL import Image, ImageDraw

detector = pipeline(
    "object-detection",
    model="facebook/detr-resnet-50"
)

colors = [
    "#FF0000", "#00FF00", "#0000FF", "#FF00FF",
    "#00FFFF", "#FFA500", "#800080", "#008000"
]

def find_obj(image):
    results = detector(image)

    draw_image = image.copy()
    draw = ImageDraw.Draw(draw_image)

    for i, item in enumerate(results):
        label = item["label"]
        score = round(item["score"] * 100, 1)
        box = item["box"]
        color = colors[i % len(colors)]

        draw.rectangle(
            [box["xmin"], box["ymin"], box["xmax"], box["ymax"]],
            outline=color,
            width=3
        )

        text = f"{label} {score}%"
        draw.rectangle(
            [box["xmin"], box["ymin"] - 20, box["xmin"] + len(text) * 7, box["ymin"]],
            fill=color
        )

        draw.text(
            (box["xmin"] + 2, box["ymin"] - 18),
            text,
            fill="white"
        )

    summary = f"총 {len(results)}개 객체 탐지됨\n\n"
    for item in results:
        summary += f"· {item['label']} ({round(item['score']*100, 1)}%)\n"

    return draw_image, summary

demo = gr.Interface(
    fn=find_obj,
    inputs=gr.Image(
        type="pil",
        label="이미지 업로드"
    ),
    outputs=[
        gr.Image(label="탐지 결과"),
        gr.Textbox(label="탐지 목록", lines=10)
    ],
    title="🔍 객체 탐지기",
    description="이미지를 업로드하면 객체를 자동으로 탐지합니다.",
)

demo.launch()  # ← share=True 제거