app.py

app.py

@@ -1,64 +1,42 @@
-# app.py
 import gradio as gr
 from transformers import AutoModelForImageClassification, AutoImageProcessor
 from PIL import Image
 import torch
 
-# ---------------- 配置模型列表 ----------------
-MODEL_LIST = [
-    "prithivMLmods/Trash-Net",
+# 三个模型
+MODELS = [
     "yangy50/garbage-classification",
-    "ahmzakif/TrashNet-Classification"
+    "ahmzakif/TrashNet-Classification",
+    "harriskr14/trashnet-vit"
 ]
 
-# ---------------- 加载模型 ----------------
-models = []
 processors = []
-loaded_model_names = []
-print("🔹 正在加载模型，请稍等...")
-
-for model_name in MODEL_LIST:
-    try:
-        processor = AutoImageProcessor.from_pretrained(model_name)
-        model = AutoModelForImageClassification.from_pretrained(model_name)
-        model.eval()
-        processors.append(processor)
-        models.append(model)
-        loaded_model_names.append(model_name)
-        print(f"✅ 加载成功: {model_name}")
-    except Exception as e:
-        print(f"❌ 加载失败: {model_name}, 错误: {e}")
+models = []
+for name in MODELS:
+    p = AutoImageProcessor.from_pretrained(name)
+    m = AutoModelForImageClassification.from_pretrained(name)
+    m.eval()
+    processors.append(p)
+    models.append(m)
 
-# ---------------- 推理函数 ----------------
 def classify_image(image: Image.Image):
     results = {}
-    for model_name, processor, model in zip(loaded_model_names, processors, models):
-        try:
-            inputs = processor(images=image, return_tensors="pt")
-            with torch.no_grad():
-                outputs = model(**inputs)
-                pred = outputs.logits.argmax(-1).item()
-                if hasattr(model.config, "id2label"):
-                    label = model.config.id2label[pred]
-                else:
-                    label = f"⚠️ 无内置 id2label，索引预测: {pred}"
-                results[model_name] = label
-        except Exception as e:
-            results[model_name] = f"❌ 预测失败: {e}"
-
-    results_text = "\n".join([f"{name}: {label}" for name, label in results.items()])
-    return results_text
+    for name, p, m in zip(MODELS, processors, models):
+        inputs = p(images=image, return_tensors="pt")
+        with torch.no_grad():
+            outputs = m(**inputs)
+            pred = outputs.logits.argmax(-1).item()
+            label = m.config.id2label.get(pred, f"id_{pred}")
+        results[name] = label
+    return "\n".join([f"{k}: {v}" for k, v in results.items()])
 
-# ---------------- Gradio 界面 ----------------
 iface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil", label="上传图片"),
-    outputs=[gr.Textbox(label="所有模型预测结果")],
-    title="垃圾分类全模型检测",
-    description="上传图片后，每个模型独立输出预测结果，不做任何人工干预。",
-    allow_flagging="never"
+    outputs="text",
+    title="三模型垃圾分类",
+    description="使用三个模型独立预测垃圾种类"
 )
 
-# ✅ 启用 API 模式（Space 可被外部调用）
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860, show_api=True)
+    iface.launch()