Create app.py

app.py

@@ -0,0 +1,147 @@
+import torch
+import gradio as gr
+from PIL import Image
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from huggingface_hub import hf_hub_download
+import clip
+
+# -------------------------------
+# 1. 台灣縣市垃圾分類規則
+# -------------------------------
+city_rules = {
+    "台北市": {
+        "一般垃圾": "丟一般垃圾袋（白色）。",
+        "紙類": "紙類須乾淨並綁好，不可含塑膠膜。",
+        "塑膠類": "乾淨塑膠可回收；髒污塑膠丟一般垃圾。",
+        "金屬類": "鐵鋁罐可回收。",
+        "玻璃類": "玻璃瓶可回收；破碎玻璃需包裹後丟一般垃圾。",
+        "食物垃圾": "分可食 / 不可食兩桶。",
+        "電池": "丟入電池回收箱。",
+        "電子產品": "需交給回收車或清潔隊。"
+    },
+    "新北市": {
+        "一般垃圾": "黑色垃圾袋。",
+        "紙類": "紙容器需清洗，紙類需綁好。",
+        "塑膃類": "PET 與其他塑膠須分開回收。",
+        "金屬類": "鐵與鋁需分類。",
+        "玻璃類": "玻璃瓶回收不分顏色。",
+        "食物垃圾": "統一廚餘桶。",
+        "電池": "請投遞電池回收箱。",
+        "電子產品": "請交由回收車或清潔隊。"
+    },
+    "高雄市": {
+        "一般垃圾": "黑色垃圾袋。",
+        "紙類": "紙容器屬紙類可回收。",
+        "塑膠類": "乾燥可回收；油污塑膠丟一般垃圾。",
+        "金屬類": "鐵鋁罐可回收。",
+        "玻璃類": "瓶類回收；破損玻璃需包裝。",
+        "食物垃圾": "各區規則不同。",
+        "電池": "投遞電池專用箱。",
+        "電子產品": "屬有害垃圾，需交回收車或清潔隊。"
+    }
+}
+
+# -------------------------------
+# 2. 載入 CLIP
+# -------------------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+clip_model, preprocess = clip.load("ViT-B/32", device=device)
+
+labels = ["一般垃圾", "紙類", "塑膠類", "金屬類", "玻璃類", "食物垃圾", "電池", "電子產品"]
+
+# -------------------------------
+# 3. LLM (phi-2)
+# -------------------------------
+tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
+llm = AutoModelForCausalLM.from_pretrained("microsoft/phi-2").to(device)
+
+def generate_reason(item_name, category):
+    prompt = f"物品「{item_name}」被歸類為「{category}」。請用一句話簡單解釋原因："
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = llm.generate(**inputs, max_new_tokens=40)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# -------------------------------
+# 4. GroundingDINO（從 HF 下載權重）
+# -------------------------------
+from groundingdino.util.inference import load_model, load_image, predict, annotate
+
+checkpoint = hf_hub_download(
+    repo_id="IDEA-Research/GroundingDINO",
+    filename="groundingdino_swint_ogc.pth"
+)
+
+config_path = hf_hub_download(
+    repo_id="IDEA-Research/GroundingDINO",
+    filename="GroundingDINO_SwinT_OGC.py"
+)
+
+dino_model = load_model(config_path, checkpoint)
+
+# -------------------------------
+# 5. 主流程：偵測 → 分類 → 解釋 → 建議
+# -------------------------------
+def classify_item(image, obj_name, city):
+    text_inputs = clip.tokenize(labels).to(device)
+    img_tensor = preprocess(image).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        logits, _ = clip_model(img_tensor, text_inputs)
+        probs = logits.softmax(dim=-1).cpu().numpy()[0]
+
+    idx = probs.argmax()
+    category = labels[idx]
+
+    reason = generate_reason(obj_name, category)
+    suggestion = city_rules[city][category]
+
+    return category, reason, suggestion
+
+
+def pipeline(image, city):
+    image.save("temp.jpg")
+
+    img_np, img_tensor = load_image("temp.jpg")
+    boxes, logits, phrases = predict(
+        model=dino_model,
+        image=img_tensor,
+        caption="bottle, can, box, plastic, cup, phone, food, battery",
+        box_threshold=0.3,
+        text_threshold=0.25
+    )
+
+    annotated = annotate(img_np, boxes, logits, phrases)
+
+    result = ""
+    for obj in phrases:
+        category, reason, suggestion = classify_item(image, obj, city)
+        result += f"""
+🧩 偵測到物品：{obj}
+📦 分類：{category}
+💡 理由：{reason}
+🚮 縣市建議（{city}）：{suggestion}
+
+"""
+
+    return annotated, result
+
+
+# -------------------------------
+# 6. Gradio 介面
+# -------------------------------
+with gr.Blocks() as demo:
+    gr.Markdown("# 🗑️ AI 垃圾分類助手（縣市版）")
+    gr.Markdown("自動偵測圖片中的垃圾 → 分類 → 產生理由 → 給縣市專屬分類建議")
+
+    with gr.Row():
+        img = gr.Image(label="上傳垃圾照片")
+        city = gr.Dropdown(["台北市", "新北市", "高雄市"], value="台北市")
+
+    btn = gr.Button("開始分析 🚀", variant="primary")
+
+    output_img = gr.Image(label="偵測結果")
+    output_text = gr.Textbox(label="分析結果", lines=12)
+
+    btn.click(pipeline, inputs=[img, city], outputs=[output_img, output_text])
+
+demo.launch()