app.py

app.py

@@ -1,41 +1,84 @@
+# app.py
 import gradio as gr
-from transformers import AutoModelForImageClassification, AutoImageProcessor
+from transformers import AutoModelForImageClassification, AutoImageProcessor, AutoFeatureExtractor
 from PIL import Image
 import torch
+from torchvision import transforms
 
-# 三个模型
-MODELS = [
+# ---------------- 配置模型 ----------------
+MODEL_LIST = [
     "yangy50/garbage-classification",
-    "ahmzakif/TrashNet-Classification",
-    "harriskr14/trashnet-vit"
+    "harriskr14/trashnet-vit",
+    "ahmzakif/TrashNet-Classification"
 ]
 
-processors = []
+# ---------------- 加载模型 ----------------
 models = []
-for name in MODELS:
-    p = AutoImageProcessor.from_pretrained(name)
-    m = AutoModelForImageClassification.from_pretrained(name)
-    m.eval()
-    processors.append(p)
-    models.append(m)
+processors = []
+loaded_model_names = []
+
+print("🔹 正在加载模型，请稍等...")
 
+for name in MODEL_LIST:
+    try:
+        model = AutoModelForImageClassification.from_pretrained(name)
+        # 尝试加载 AutoImageProcessor
+        try:
+            processor = AutoImageProcessor.from_pretrained(name)
+        except OSError:
+            # 尝试 AutoFeatureExtractor
+            try:
+                processor = AutoFeatureExtractor.from_pretrained(name)
+            except OSError:
+                processor = None
+                print(f"⚠️ 模型 {name} 没有自带处理器，将使用默认 transforms")
+        model.eval()
+        models.append(model)
+        processors.append(processor)
+        loaded_model_names.append(name)
+        print(f"✅ 加载成功: {name}")
+    except Exception as e:
+        print(f"❌ 加载失败: {name}, 错误: {e}")
+
+# ---------------- 默认图片处理 ----------------
+default_preprocess = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])
+])
+
+# ---------------- 推理函数 ----------------
 def classify_image(image: Image.Image):
     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()])
+    for name, model, processor in zip(loaded_model_names, models, processors):
+        try:
+            if processor:
+                inputs = processor(images=image, return_tensors="pt")
+            else:
+                inputs = default_preprocess(image).unsqueeze(0)
+            with torch.no_grad():
+                outputs = model(**inputs)
+                pred_idx = outputs.logits.argmax(-1).item()
+                if hasattr(model.config, "id2label"):
+                    label = model.config.id2label[pred_idx]
+                else:
+                    label = f"索引预测: {pred_idx}"
+                results[name] = label
+        except Exception as e:
+            results[name] = f"❌ 预测失败: {e}"
+
+    # 格式化输出
+    results_text = "\n".join([f"{name}: {label}" for name, label in results.items()])
+    return results_text
 
+# ---------------- Gradio 界面 ----------------
 iface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil", label="上传图片"),
-    outputs="text",
-    title="三模型垃圾分类",
-    description="使用三个模型独立预测垃圾种类"
+    outputs=[gr.Textbox(label="模型预测结果")],
+    title="垃圾分类模型检测",
+    description="上传图片，每个模型独立预测，输出所有结果。"
 )
 
 if __name__ == "__main__":