app.py

app.py

@@ -1,57 +1,41 @@
-# app.py
 import gradio as gr
 from transformers import AutoModelForImageClassification, AutoImageProcessor
 from PIL import Image
 import torch
 
-# ---------------- 模型列表 ----------------
-MODEL_LIST = [
+# 三个模型
+MODELS = [
     "yangy50/garbage-classification",
-    "harriskr14/trashnet-vit",
-    "prithivMLmods/Trash-Net"
+    "ahmzakif/TrashNet-Classification",
+    "harriskr14/trashnet-vit"
 ]
 
-# ---------------- 加载模型 ----------------
-models = []
 processors = []
-loaded_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_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 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()
-                label = model.config.id2label.get(pred, str(pred))
-                results.append(f"{name}: {label}")
-        except Exception as e:
-            results.append(f"{name}: ❌ 预测失败 ({e})")
-
-    return "\n".join(results)
+    results = {}
+    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="模型预测结果", lines=6),
-    title="多模型垃圾分类",
-    description="使用以下模型进行独立预测：yangy50、harriskr14、prithivMLmods。"
+    outputs="text",
+    title="三模型垃圾分类",
+    description="使用三个模型独立预测垃圾种类"
 )
 
 if __name__ == "__main__":