app.py

app.py

@@ -5,19 +5,23 @@ from PIL import Image
 import torch
 from collections import Counter
 
-# ---------------- 模型列表 ----------------
 MODEL_LIST = [
     "prithivMLmods/Trash-Net",
     "yangy50/garbage-classification",
     "eunoiawiira-vgg-realwaste-classification"
 ]
 
-# ---------------- 加载模型 ----------------
+PRIORITY_ORDER = [
+    "yangy50/garbage-classification",
+    "eunoiawiira-vgg-realwaste-classification",
+    "prithivMLmods/Trash-Net"
+]
+
 models = []
 processors = []
 devices = []
 
-print("🔹 正在加载模型，请稍等...")
+print("Loading models...")
 
 for model_name in MODEL_LIST:
     try:
@@ -31,11 +35,10 @@ for model_name in MODEL_LIST:
         processors.append(processor)
         models.append(model)
         devices.append(next(model.parameters()).device)
-        print(f"✅ 加载完成: {model_name}")
+        print(f"Loaded: {model_name}")
     except Exception as e:
-        print(f"❌ 加载失败: {model_name}, 错误: {e}")
+        print(f"Failed to load {model_name}, error: {e}")
 
-# ---------------- 推理函数 ----------------
 def classify_image(image: Image.Image):
     results = {}
     for model_name, processor, model, device in zip(MODEL_LIST, processors, models, devices):
@@ -52,24 +55,33 @@ def classify_image(image: Image.Image):
 
     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}"
+    label_counts = Counter(valid_labels)
+
+    if len(label_counts) == 0:
+        final_label = "Unknown"
+    elif len(label_counts) == 1:
+        final_label = valid_labels[0]
+    else:
+        for model_name in PRIORITY_ORDER:
+            if model_name in results and not results[model_name].startswith("error"):
+                final_label = results[model_name]
+                break
+
+    results_text += f"\n\nFinal Label: {final_label}"
     return results_text
 
-# ---------------- Gradio 界面 ----------------
 iface = gr.Interface(
     fn=classify_image,
-    inputs=gr.Image(type="pil", label="上传图片"),
-    outputs=[gr.Textbox(label="模型预测结果")],
-    title="垃圾分类多模型检测",
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=[gr.Textbox(label="Model Predictions")],
+    title="Multi-Model Trash Classification",
     description=(
-        "上传图片后，使用以下模型进行垃圾分类，每个模型结果单独输出：\n"
-        "1. prithivMLmods/Trash-Net\n"
-        "2. yangy50/garbage-classification\n"
-        "3. eunoiawiira-vgg-realwaste-classification\n"
-        "最终标签通过投票法生成"
+        "Uploads an image and classifies trash using three models.\n"
+        "All model predictions are displayed and the final label is selected by priority:\n"
+        "1. yangy50/garbage-classification\n"
+        "2. eunoiawiira-vgg-realwaste-classification\n"
+        "3. prithivMLmods/Trash-Net"
     )
 )