# 사용자 이미지 업로드 -> 이미지, 텍스트로 결과 출력
import gradio as gr
from ultralytics import YOLO
import numpy as np
from PIL import Image
import yaml

# 모델 로드
model = YOLO("best.pt")

# 클래스 이름 불러오기
with open("data.yaml", "r") as f:
    data = yaml.safe_load(f)
    class_names = data["names"]

def detect_and_summarize(image):
    # PIL → numpy
    img_array = np.array(image)

    # YOLO 추론
    results = model(img_array)[0]

    # 바운딩 박스 그리기 + 결과 클래스 추출
    detected_classes = []
    for box in results.boxes:
        cls_id = int(box.cls[0].item())
        class_name = class_names[cls_id]
        detected_classes.append(class_name)

    # 중복 제거
    detected_classes = list(set(detected_classes))
    summary = ", ".join(detected_classes) if detected_classes else "탐지된 객체 없음"

    # 시각화된 이미지 얻기
    plotted_img = results.plot()  # returns numpy array
    result_image = Image.fromarray(plotted_img)

    return result_image, f"탐지된 클래스: {summary}"

# Gradio 인터페이스
with gr.Blocks() as demo:
    gr.Markdown("## 🍅 YOLO 객체 탐지 결과 템플릿 1")
    gr.Markdown("이미지를 업로드하면 YOLO 모델이 객체를 탐지하고, 결과를 요약해줍니다.")

    input_image = gr.Image(type="pil", label="이미지 업로드")
    submit_btn = gr.Button("탐지 시작")
    output_image = gr.Image(label="탐지 결과 이미지")
    output_text = gr.Textbox(label="탐지된 클래스 요약")

    submit_btn.click(fn=detect_and_summarize, inputs=input_image, outputs=[output_image, output_text])

demo.launch()