# tab 1 : 사용자가 test 데이터 셋에서 원하는 번호를 선택하여 결과 확인
# tabl 2 : 사용자가 이미지 업로드
import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np
import yaml
import os
import zipfile


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

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

# test.zip 파일 압축 해제
with zipfile.ZipFile("test.zip", 'r') as zip_ref:
    zip_ref.extractall("test")  # 폴더명은 원하는 경로로 지정

# 객체 탐지 함수
def detect_objects(image):
    results = model(image)[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)

    result_image = Image.fromarray(results.plot())
    summary = ", ".join(set(detected_classes)) if detected_classes else "탐지된 객체 없음"

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

# 이미지 파일 리스트 가져오기
def get_image_file_list():
    image_dir = "test/images"
    file_list = sorted([
        f for f in os.listdir(image_dir) if f.endswith((".jpg", ".png"))
    ])
    return file_list

# 선택된 인덱스로 이미지 로드
def load_image_by_index(index_display: int):
    index = index_display - 1  # 사용자에게는 1부터 보여줬기 때문에
    file_list = get_image_file_list()
    if 0 <= index < len(file_list):
        image_path = os.path.join("test/images", file_list[index])
        return Image.open(image_path).convert("RGB")
    else:
        return Image.new("RGB", (300, 300), color="gray")

# Gradio 앱 시작
with gr.Blocks() as demo:
    gr.Markdown("## ✌ YOLO 객체 탐지 템플릿 2 - 탭 기반 멀티기능 데모")
    gr.Markdown("Tab 1: 테스트 이미지 선택 / Tab 2: 사용자 이미지 업로드")

    with gr.Tab("📁 테스트 데이터셋에서 실행"):
        file_list = get_image_file_list()
        index_choices = [i + 1 for i in range(len(file_list))]  # 1부터 시작

        dropdown = gr.Dropdown(choices=index_choices, value=1, label="테스트 이미지 번호 선택")
        load_btn = gr.Button("이미지 불러오기")
        selected_image = gr.Image(label="선택된 이미지", type="pil")
        detect_btn = gr.Button("탐지 실행")
        result_img = gr.Image(label="탐지 결과 이미지")
        result_txt = gr.Textbox(label="탐지된 클래스 요약")

        load_btn.click(fn=load_image_by_index, inputs=dropdown, outputs=selected_image)
        detect_btn.click(fn=detect_objects, inputs=selected_image, outputs=[result_img, result_txt])

    with gr.Tab("🖼️ 사용자 이미지 업로드"):
        user_img = gr.Image(type="pil", label="이미지 업로드")
        user_btn = gr.Button("탐지 실행")
        user_result_img = gr.Image(label="탐지 결과 이미지")
        user_result_txt = gr.Textbox(label="탐지된 클래스 요약")
        user_btn.click(fn=detect_objects, inputs=user_img, outputs=[user_result_img, user_result_txt])

demo.launch()