import os
from transformers import pipeline
from PIL import Image
import gradio as gr

# 🔹 객체 탐지 모델 로딩
detector = pipeline(
    task="zero-shot-object-detection",
    model="google/owlv2-base-patch16-ensemble"
)

# 🔹 탐지할 재료 라벨 리스트
candidate_labels = [
    "salmon", "chicken breast", "broccoli", "lettuce", "mushroom",
    "bell pepper", "onion", "cherry tomato", "egg", "milk", "cheese", "garlic"
]

# 🔹 영어 → 한글 라벨 변환 딕셔너리
label_ko = {
    "salmon": "연어",
    "chicken breast": "닭가슴살",
    "broccoli": "브로콜리",
    "lettuce": "상추",
    "mushroom": "버섯",
    "bell pepper": "피망",
    "onion": "양파",
    "cherry tomato": "방울토마토",
    "egg": "계란",
    "milk": "우유",
    "cheese": "치즈",
    "garlic": "마늘"
}

# 🔹 재료 탐지 함수
def detect_ingredients(image):
    if image is None:
        return "이미지를 먼저 업로드해주세요."

    outputs = detector(image, candidate_labels=candidate_labels, threshold=0.2)
    detected_labels = list(set([o["label"] for o in outputs]))
    
    if not detected_labels:
        return "재료를 인식하지 못했습니다. 더 명확한 이미지를 사용해주세요."
    
    # 영어 → 한글 변환
    translated_labels = [label_ko.get(label, label) for label in detected_labels]
    return ", ".join(translated_labels)

# 🔹 Gradio UI 구성
with gr.Blocks() as demo:
    gr.Markdown("## 🧊 냉장고 재료 탐지 시스템")

    image_input = gr.Image(type="pil", label="냉장고 이미지 업로드")
    detect_button = gr.Button("재료 탐지하기")
    ingredient_output = gr.Textbox(label="📌 탐지된 재료 (한글)")

    detect_button.click(fn=detect_ingredients, inputs=image_input, outputs=ingredient_output)

if __name__ == "__main__":
    demo.launch()