app.py

app.py

@@ -8,13 +8,7 @@ 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"
+    "eunoiawiira/vgg-realwaste-classification"
 ]
 
 models = []
@@ -22,7 +16,6 @@ processors = []
 devices = []
 
 print("Loading models...")
-
 for model_name in MODEL_LIST:
     try:
         processor = AutoImageProcessor.from_pretrained(model_name)
@@ -37,7 +30,7 @@ for model_name in MODEL_LIST:
         devices.append(next(model.parameters()).device)
         print(f"Loaded: {model_name}")
     except Exception as e:
-        print(f"Failed to load {model_name}, error: {e}")
+        print(f"Failed to load {model_name}: {e}")
 
 def classify_image(image: Image.Image):
     results = {}
@@ -51,24 +44,15 @@ def classify_image(image: Image.Image):
             label = model.config.id2label[pred]
             results[model_name] = label
         except Exception as e:
-            results[model_name] = f"error: {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")]
-    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}"
+    final_label = Counter(valid_labels).most_common(1)[0][0] if valid_labels else "Unknown"
+    results_text += f"\n\nFinal Label (voting): {final_label}"
     return results_text
 
 iface = gr.Interface(
@@ -77,11 +61,11 @@ iface = gr.Interface(
     outputs=[gr.Textbox(label="Model Predictions")],
     title="Multi-Model Trash Classification",
     description=(
-        "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"
+        "Upload an image, and the following models will classify it:\n"
+        "1. prithivMLmods/Trash-Net\n"
+        "2. yangy50/garbage-classification\n"
+        "3. eunoiawiira/vgg-realwaste-classification\n"
+        "The final label is determined by majority vote."
     )
 )