Create app.py

app.py

@@ -0,0 +1,107 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from PIL import Image
+import json
+from sentence_transformers import SentenceTransformer, util
+from openai import OpenAI
+import os
+
+# ========== 1. Load Image Classification Model ==========
+class MobileNetClassifier(nn.Module):
+    def __init__(self, num_classes=12):
+        super().__init__()
+        self.model = torch.hub.load("pytorch/vision:v0.10.0", "mobilenet_v2", pretrained=False)
+        self.model.classifier[1] = nn.Linear(1280, num_classes)
+
+    def forward(self, x):
+        return self.model(x)
+
+# Load model
+device = "cpu"
+model = MobileNetClassifier(num_classes=12).to(device)
+model.load_state_dict(torch.load("mobilenet_trash.pth", map_location=device))
+model.eval()
+
+labels = [
+    "cardboard", "glass", "metal", "paper", "plastic", 
+    "trash", "battery", "shoes", "clothes", "green-glass",
+    "brown-glass", "white-glass"
+]
+
+transform = transforms.Compose([
+    transforms.Resize((224,224)),
+    transforms.ToTensor()
+])
+
+# ========== 2. Load QAS + Recycle Database ==========
+qas = json.load(open("qas.json", "r", encoding="utf-8"))
+recycle = json.load(open("recycle_data.json", "r", encoding="utf-8"))
+recycle_dict = {item["name"]: item for item in recycle}
+
+# Embedding Model
+embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+qas_questions = [item["question"] for item in qas]
+qas_embeddings = embedder.encode(qas_questions, convert_to_tensor=True)
+
+# ========== 3. Helper Functions ==========
+def classify_image(image):
+    img = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        pred = model(img)[0]
+    idx = torch.argmax(pred).item()
+    return labels[idx]
+
+def search_recycle(label):
+    if label in recycle_dict:
+        item = recycle_dict[label]
+        return f"♻ {label}\n\n回收方式：{item['notes']}"
+    return f"⚠ 找不到「{label}」的回收資料，請依一般原則：可分離材質可回收。"
+
+def rag_question(query):
+    emb = embedder.encode(query, convert_to_tensor=True)
+    scores = util.cos_sim(emb, qas_embeddings)[0]
+    idx = torch.argmax(scores).item()
+    score = scores[idx].item()
+
+    if score > 0.7:
+        return qas[idx]["answer"]
+    return None
+
+client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
+
+def llm_fallback(query):
+    msg = f"你是台灣垃圾分類助理，請回答：{query}"
+    resp = client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[{"role": "user", "content": msg}]
+    )
+    return resp.choices[0].message["content"]
+
+# ========== 4. Master Function ==========
+def assistant(text, image):
+    if image is not None:
+        image = Image.fromarray(image)
+        label = classify_image(image)
+        result = search_recycle(label)
+        return f"🔍 辨識結果：{label}\n\n{result}"
+
+    if text:
+        rag_answer = rag_question(text)
+        if rag_answer:
+            return rag_answer
+        return llm_fallback(text)
+
+    return "請輸入問題或上傳圖片。"
+
+# ========== 5. Gradio UI ==========
+ui = gr.Interface(
+    fn=assistant,
+    inputs=[gr.Textbox(label="輸入問題"), gr.Image(type="numpy")],
+    outputs="text",
+    title="垃圾分類智慧助理",
+    description="上傳圖片或輸入問題，協助你判斷台灣垃圾分類方式"
+)
+
+ui.launch()