app.py

app.py

@@ -3,8 +3,9 @@ import gradio as gr
 from transformers import AutoModelForImageClassification, AutoImageProcessor
 from PIL import Image
 import torch
+from collections import Counter
 
-# ---------------- 配置模型列表，包含你提到的全部模型 ----------------
+# ---------------- 模型列表 ----------------
 MODEL_LIST = [
     "prithivMLmods/Trash-Net",
     "yangy50/garbage-classification",
@@ -14,14 +15,21 @@ MODEL_LIST = [
 # ---------------- 加载模型 ----------------
 models = []
 processors = []
+devices = []
 print("🔹 正在加载模型，请稍等...")
+
 for model_name in MODEL_LIST:
     try:
         processor = AutoImageProcessor.from_pretrained(model_name)
-        model = AutoModelForImageClassification.from_pretrained(model_name)
+        model = AutoModelForImageClassification.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto" if torch.cuda.is_available() else None
+        )
         model.eval()
         processors.append(processor)
         models.append(model)
+        devices.append(next(model.parameters()).device)
         print(f"✅ 加载完成: {model_name}")
     except Exception as e:
         print(f"❌ 加载失败: {model_name}, 错误: {e}")
@@ -29,30 +37,41 @@ for model_name in MODEL_LIST:
 # ---------------- 推理函数 ----------------
 def classify_image(image: Image.Image):
     results = {}
-    # 遍历所有模型进行预测
-    for model_name, processor, model in zip(MODEL_LIST, processors, models):
-        inputs = processor(images=image, return_tensors="pt")
-        with torch.no_grad():
-            outputs = model(**inputs)
+    for model_name, processor, model, device in zip(MODEL_LIST, processors, models, devices):
+        try:
+            inputs = processor(images=image, return_tensors="pt")
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+            with torch.no_grad():
+                outputs = model(**inputs)
             pred = outputs.logits.argmax(-1).item()
             label = model.config.id2label[pred]
             results[model_name] = label
+        except Exception as e:
+            results[model_name] = f"error: {e}"
 
-    # 格式化输出，每个模型的结果单独显示
+    # 格式化输出
     results_text = "\n".join([f"{name}: {label}" for name, label in results.items()])
 
+    # 计算最终标签（投票法）
+    valid_labels = [lbl for lbl in results.values() if not lbl.startswith("error")]
+    final_label = Counter(valid_labels).most_common(1)[0][0] if valid_labels else "Unknown"
+
+    results_text += f"\n\n最终标签: {final_label}"
     return results_text
 
 # ---------------- Gradio 界面 ----------------
 iface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil", label="上传图片"),
-    outputs=[gr.Textbox(label="所有模型预测结果")],
+    outputs=[gr.Textbox(label="模型预测结果")],
     title="垃圾分类多模型检测",
-    description="上传图片后，使用以下模型进行垃圾分类，每个模型结果单独输出：\n"
-                "1. prithivMLmods/Trash-Net\n"
-                "2. yangy50/garbage-classification\n"
-                "3. eunoiawiira-vgg-realwaste-classification\n"
+    description=(
+        "上传图片后，使用以下模型进行垃圾分类，每个模型结果单独输出：\n"
+        "1. prithivMLmods/Trash-Net\n"
+        "2. yangy50/garbage-classification\n"
+        "3. eunoiawiira-vgg-realwaste-classification\n"
+        "最终标签通过投票法生成"
+    )
 )
 
 if __name__ == "__main__":