app.py

app.py

@@ -4,65 +4,54 @@ from transformers import AutoModelForImageClassification, AutoImageProcessor
 from PIL import Image
 import torch
 
-# ---------------- 配置模型列表 ----------------
+# ---------------- 模型列表 ----------------
 MODEL_LIST = [
-    "prithivMLmods/Trash-Net",
     "yangy50/garbage-classification",
-    "eunoiawiira-vgg-realwaste-classification",
-    "ahmzakif/TrashNet-Classification",
-    "ee8225-group4-vit-trashnet-enhanced",
-    "harriskr14/trashnet-vit"
+    "harriskr14/trashnet-vit",
+    "prithivMLmods/Trash-Net"
 ]
 
 # ---------------- 加载模型 ----------------
 models = []
 processors = []
-loaded_model_names = []
-print("🔹 正在加载模型，请稍等...")
+loaded_names = []
 
+print("🔹 正在加载模型，请稍等...")
 for model_name in MODEL_LIST:
     try:
         processor = AutoImageProcessor.from_pretrained(model_name)
         model = AutoModelForImageClassification.from_pretrained(model_name)
         model.eval()
-        # 只加载模型自带 id2label，不手动干预
-        if not hasattr(model.config, "id2label"):
-            print(f"⚠️ 模型 {model_name} 没有内置 id2label，预测可能失败")
         processors.append(processor)
         models.append(model)
-        loaded_model_names.append(model_name)
+        loaded_names.append(model_name)
         print(f"✅ 加载成功: {model_name}")
     except Exception as e:
         print(f"❌ 加载失败: {model_name}, 错误: {e}")
 
 # ---------------- 推理函数 ----------------
 def classify_image(image: Image.Image):
-    results = {}
-    for model_name, processor, model in zip(loaded_model_names, processors, models):
+    results = []
+    for name, processor, model in zip(loaded_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
+                label = model.config.id2label.get(pred, str(pred))
+                results.append(f"{name}: {label}")
         except Exception as e:
-            results[model_name] = f"❌ 预测失败: {e}"
+            results.append(f"{name}: ❌ 预测失败 ({e})")
 
-    # 格式化输出
-    results_text = "\n".join([f"{name}: {label}" for name, label in results.items()])
-    return results_text
+    return "\n".join(results)
 
 # ---------------- Gradio 界面 ----------------
 iface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil", label="上传图片"),
-    outputs=[gr.Textbox(label="所有模型预测结果")],
-    title="垃圾分类全模型检测",
-    description="上传图片后，每个模型独立输出预测结果，不做任何人工干预。"
+    outputs=gr.Textbox(label="模型预测结果", lines=6),
+    title="多模型垃圾分类",
+    description="使用以下模型进行独立预测：yangy50、harriskr14、prithivMLmods。"
 )
 
 if __name__ == "__main__":