app.py

app.py

@@ -1,15 +1,30 @@
 # app.py
-import gradio as gr
-from transformers import AutoModelForImageClassification, AutoImageProcessor, AutoFeatureExtractor
+import subprocess
+import sys
+
+# ---------------- 自动安装缺失依赖 ----------------
+def install(package):
+    subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+
+for pkg in ["torch", "torchvision", "transformers", "gradio", "Pillow"]:
+    try:
+        __import__(pkg if pkg != "Pillow" else "PIL")
+    except ModuleNotFoundError:
+        print(f"⚠️ 未找到 {pkg}，正在自动安装...")
+        install(pkg)
+        print(f"✅ {pkg} 安装完成")
+
+# ---------------- 导入库 ----------------
+from transformers import AutoModelForImageClassification, AutoImageProcessor
 from PIL import Image
 import torch
-from torchvision import transforms
+import gradio as gr
 
-# ---------------- 配置模型 ----------------
+# ---------------- 模型列表 ----------------
 MODEL_LIST = [
     "yangy50/garbage-classification",
-    "harriskr14/trashnet-vit",
-    "ahmzakif/TrashNet-Classification"
+    "ahmzakif/TrashNet-Classification",
+    "harriskr14/trashnet-vit"
 ]
 
 # ---------------- 加载模型 ----------------
@@ -18,55 +33,34 @@ processors = []
 loaded_model_names = []
 
 print("🔹 正在加载模型，请稍等...")
-
-for name in MODEL_LIST:
+for model_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")
+        processor = AutoImageProcessor.from_pretrained(model_name)
+        model = AutoModelForImageClassification.from_pretrained(model_name)
         model.eval()
-        models.append(model)
         processors.append(processor)
-        loaded_model_names.append(name)
-        print(f"✅ 加载成功: {name}")
+        models.append(model)
+        loaded_model_names.append(model_name)
+        print(f"✅ 加载成功: {model_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])
-])
+        print(f"❌ 加载失败: {model_name}, 错误: {e}")
 
 # ---------------- 推理函数 ----------------
 def classify_image(image: Image.Image):
     results = {}
-    for name, model, processor in zip(loaded_model_names, models, processors):
+    for model_name, processor, model in zip(loaded_model_names, processors, models):
         try:
-            if processor:
-                inputs = processor(images=image, return_tensors="pt")
-            else:
-                inputs = default_preprocess(image).unsqueeze(0)
+            inputs = processor(images=image, return_tensors="pt")
             with torch.no_grad():
                 outputs = model(**inputs)
-                pred_idx = outputs.logits.argmax(-1).item()
+                pred = outputs.logits.argmax(-1).item()
                 if hasattr(model.config, "id2label"):
-                    label = model.config.id2label[pred_idx]
+                    label = model.config.id2label[pred]
                 else:
-                    label = f"索引预测: {pred_idx}"
-                results[name] = label
+                    label = f"⚠️ 无内置 id2label，索引预测: {pred}"
+                results[model_name] = label
         except Exception as e:
-            results[name] = f"❌ 预测失败: {e}"
+            results[model_name] = f"❌ 预测失败: {e}"
 
     # 格式化输出
     results_text = "\n".join([f"{name}: {label}" for name, label in results.items()])
@@ -76,9 +70,9 @@ def classify_image(image: Image.Image):
 iface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil", label="上传图片"),
-    outputs=[gr.Textbox(label="模型预测结果")],
+    outputs=[gr.Textbox(label="所有模型预测结果")],
     title="垃圾分类模型检测",
-    description="上传图片，每个模型独立预测，输出所有结果。"
+    description="上传图片后，每个模型独立输出预测结果，不做任何人工干预。"
 )
 
 if __name__ == "__main__":