app.py

app.py

@@ -1,41 +1,64 @@
+# app.py
 import gradio as gr
 from transformers import AutoModelForImageClassification, AutoImageProcessor
 from PIL import Image
 import torch
 
-# 三个模型
-MODELS = [
+# ---------------- 配置模型列表，包含你提到的全部模型 ----------------
+MODEL_LIST = [
+    "prithivMLmods/Trash-Net",
     "yangy50/garbage-classification",
+    "eunoiawiira-vgg-realwaste-classification",
     "ahmzakif/TrashNet-Classification",
+    "ee8225-group4-vit-trashnet-enhanced",
     "harriskr14/trashnet-vit"
 ]
 
-processors = []
+# ---------------- 加载模型 ----------------
 models = []
-for name in MODELS:
-    p = AutoImageProcessor.from_pretrained(name)
-    m = AutoModelForImageClassification.from_pretrained(name)
-    m.eval()
-    processors.append(p)
-    models.append(m)
+processors = []
+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)
+        print(f"✅ 加载完成: {model_name}")
+    except Exception as e:
+        print(f"❌ 加载失败: {model_name}, 错误: {e}")
 
+# ---------------- 推理函数 ----------------
 def classify_image(image: Image.Image):
     results = {}
-    for name, p, m in zip(MODELS, processors, models):
-        inputs = p(images=image, return_tensors="pt")
+    # 遍历所有模型进行预测
+    for model_name, processor, model in zip(MODEL_LIST, processors, models):
+        inputs = processor(images=image, return_tensors="pt")
         with torch.no_grad():
-            outputs = m(**inputs)
+            outputs = model(**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()])
+            label = model.config.id2label[pred]
+            results[model_name] = label
+
+    # 格式化输出，每个模型的结果单独显示
+    results_text = "\n".join([f"{name}: {label}" for name, label in results.items()])
+
+    return results_text
 
+# ---------------- Gradio 界面 ----------------
 iface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil", label="上传图片"),
-    outputs="text",
-    title="三模型垃圾分类",
-    description="使用三个模型独立预测垃圾种类"
+    outputs=[gr.Textbox(label="所有模型预测结果")],
+    title="垃圾分类多模型检测",
+    description="上传图片后，使用以下模型进行垃圾分类，每个模型结果单独输出：\n"
+                "1. prithivMLmods/Trash-Net\n"
+                "2. yangy50/garbage-classification\n"
+                "3. eunoiawiira-vgg-realwaste-classification\n"
+                "4. ahmzakif/TrashNet-Classification\n"
+                "5. ee8225-group4-vit-trashnet-enhanced\n"
+                "6. harriskr14/trashnet-vit"
 )
 
 if __name__ == "__main__":