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

classifier = pipeline(
    task="zero-shot-image-classification",
    model="openai/clip-vit-base-patch32"
)

LABELS = [
    "cardboard waste",
    "glass waste",
    "metal waste",
    "paper waste",
    "plastic waste",
    "general trash"
]

CLASS_ZH = {
    "cardboard waste": "紙板／紙箱類",
    "glass waste": "玻璃類",
    "metal waste": "金屬類",
    "paper waste": "紙類",
    "plastic waste": "塑膠類",
    "general trash": "一般垃圾"
}

RECYCLE_TIPS = {
    "cardboard waste": "建議壓扁後回收，若沾有大量油污或食物殘渣，應依當地規定處理。",
    "glass waste": "建議清空內容物後回收，破玻璃需妥善包裝，避免割傷清潔人員。",
    "metal waste": "建議清空內容物後回收，鋁罐、鐵罐通常可歸入金屬回收。",
    "paper waste": "乾淨紙類可回收，若嚴重沾油、沾水或污染，可能需作一般垃圾處理。",
    "plastic waste": "建議清空並簡單沖洗後回收，依塑膠材質與當地規則分類。",
    "general trash": "此類較可能為一般垃圾，建議確認是否仍有可回收部分。"
}

def predict_garbage(image):
    if image is None:
        return None, "請先上傳一張圖片。"

    results = classifier(image, candidate_labels=LABELS)

    best = results[0]
    label = best["label"]
    score = best["score"]

    chinese_name = CLASS_ZH.get(label, label)
    tip = RECYCLE_TIPS.get(label, "請依照當地垃圾分類規則處理。")

    top_text = ""

    for item in results:
        item_label = item["label"]
        item_score = item["score"]
        zh = CLASS_ZH.get(item_label, item_label)
        top_text += f"{item_label}（{zh}）：{item_score * 100:.2f}%\n"

    output_text = f"""預測結果：{label}

中文類別：{chinese_name}

信心分數：{score:.4f}

信心百分比：{score * 100:.2f}%

分類建議：
{tip}

各類別預測結果：
{top_text}
"""

    return image, output_text

custom_css = """
.gradio-container {
    max-width: 1100px !important;
    margin: auto !important;
}

#title-block {
    text-align: center;
    padding: 22px 12px 10px 12px;
}

#title-block h1 {
    font-size: 34px;
    margin-bottom: 8px;
}

#title-block p {
    font-size: 16px;
    opacity: 0.85;
}
"""

with gr.Blocks(css=custom_css, title="AI 垃圾分類影像辨識系統") as demo:
    gr.HTML("""
    <div id="title-block">
        <h1>AI 垃圾分類影像辨識系統</h1>
        <p>上傳垃圾圖片，系統會辨識 cardboard、glass、metal、paper、plastic、trash 類別</p>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="pil", label="上傳垃圾圖片")
            submit_btn = gr.Button("開始辨識", variant="primary")
            clear_btn = gr.ClearButton([input_image], value="清除圖片")

        with gr.Column(scale=1):
            output_image = gr.Image(type="pil", label="輸入圖片預覽")
            output_text = gr.Textbox(label="模型預測結果", lines=14)

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

    gr.Markdown("""
### 使用說明

1. 上傳一張垃圾圖片。
2. 點選「開始辨識」。
3. 系統會輸出垃圾類別、中文說明、信心分數與分類建議。
""")

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