Update app.py

app.py

@@ -1,45 +1,130 @@
 import gradio as gr
-from ultralytics import YOLO
+from transformers import pipeline
+from PIL import Image
 
-model = YOLO("best.pt")
+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": "一般垃圾"
+}
 
-class_zh = {
-    "cardboard": "紙板 / 紙箱類",
-    "glass": "玻璃類",
-    "metal": "金屬類",
-    "paper": "紙類",
-    "plastic": "塑膠類",
-    "trash": "一般垃圾"
+RECYCLE_TIPS = {
+    "cardboard waste": "建議壓扁後回收，若沾有大量油污或食物殘渣，應依當地規定處理。",
+    "glass waste": "建議清空內容物後回收，破玻璃需妥善包裝，避免割傷清潔人員。",
+    "metal waste": "建議清空內容物後回收，鋁罐、鐵罐通常可歸入金屬回收。",
+    "paper waste": "乾淨紙類可回收，若嚴重沾油、沾水或污染，可能需作一般垃圾處理。",
+    "plastic waste": "建議清空並簡單沖洗後回收，依塑膠材質與當地規則分類。",
+    "general trash": "此類較可能為一般垃圾，建議確認是否仍有可回收部分。"
 }
 
 def predict_garbage(image):
-    results = model.predict(image, verbose=False)
-    result = results[0]
+    if image is None:
+        return None, "請先上傳一張圖片。"
+
+    results = classifier(image, candidate_labels=LABELS)
+
+    best = results[0]
+    label = best["label"]
+    score = best["score"]
 
-    class_id = result.probs.top1
-    confidence = result.probs.top1conf.item()
-    class_name = result.names[class_id]
+    chinese_name = CLASS_ZH.get(label, label)
+    tip = RECYCLE_TIPS.get(label, "請依照當地垃圾分類規則處理。")
 
-    chinese_name = class_zh.get(class_name, class_name)
+    top_text = ""
 
-    output_text = f"""
-預測類別：{class_name}
+    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"
 
-中文說明：{chinese_name}
+    output_text = f"""預測結果：{label}
 
-信心分數：{confidence:.4f}
+中文類別：{chinese_name}
 
-信心百分比：{confidence * 100:.2f}%
+信心分數：{score:.4f}
+
+信心百分比：{score * 100:.2f}%
+
+分類建議：
+{tip}
+
+各類別預測結果：
+{top_text}
 """
 
-    return output_text
+    return image, output_text
 
-demo = gr.Interface(
-    fn=predict_garbage,
-    inputs=gr.Image(type="pil", label="上傳垃圾圖片"),
-    outputs=gr.Textbox(label="模型預測結果"),
-    title="AI 垃圾分類影像辨識系統",
-    description="使用 YOLOv8 模型辨識垃圾圖片類別，包含 cardboard、glass、metal、paper、plastic、trash 六種類別。"
-)
+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. 系統會輸出垃圾類別、中文說明、信心分數與分類建議。
+""")
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()