From 8a2ea1abd7b2aaf42c25c80dfef97c2bf35054ff Mon Sep 17 00:00:00 2001
From: zhenghu <1831829219@qq.com>
Date: Wed, 15 Apr 2026 15:09:17 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BD=BF=E7=94=A8=20ViT=20=E4=B8=93?=
 =?UTF-8?q?=E7=94=A8=E7=B1=BB=E5=8A=A0=E8=BD=BD=E6=A8=A1=E5=9E=8B=EF=BC=8C?=
 =?UTF-8?q?=E7=A7=BB=E9=99=A4=20Auto=20=E7=B1=BB=E5=85=BC=E5=AE=B9?=
 =?UTF-8?q?=E5=9B=9E=E9=80=80=E9=80=BB=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

  - 将 AutoImageProcessor / AutoModelForImageClassification 替换为
    ViTImageProcessor / ViTForImageClassification
  - 移除 ValueError 异常捕获回退逻辑，改为直接使用 ViT 类加载，
    避免因 preprocessor_config.json 中旧版 image_processor_type
    导致 Auto 类无法识别的问题
---
 app.py | 24 ++++++++----------------
 1 file changed, 8 insertions(+), 16 deletions(-)

diff --git a/app.py b/app.py
index 16a59b4..b1359ae 100644
--- a/app.py
+++ b/app.py
@@ -12,7 +12,7 @@ from PIL import Image
 # 国内 HuggingFace 镜像加速
 os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")
 
-from transformers import AutoImageProcessor, AutoModelForImageClassification, pipeline  # noqa: E402
+from transformers import ViTForImageClassification, ViTImageProcessor, pipeline  # noqa: E402
 
 # ─── Disease Label Mapping ──────────────────────────────────────────────────
 # 模型输出 LABEL_0 ~ LABEL_6，映射为实际病害名称
@@ -91,21 +91,13 @@ html, body, [class*="css"] {
 def load_model():
     """加载 HuggingFace 模型（首次运行自动下载，约 343MB）"""
     model_name = "Dmitry43243242/banana-disease-leaf-model"
-    try:
-        classifier = pipeline(
-            "image-classification",
-            model=model_name,
-        )
-    except ValueError:
-        # 模型的 preprocessor_config.json 中 image_processor_type 可能是旧版名称
-        # 手动加载并构造 pipeline
-        processor = AutoImageProcessor.from_pretrained(
-            model_name, trust_remote_code=False
-        )
-        model = AutoModelForImageClassification.from_pretrained(model_name)
-        classifier = pipeline(
-            "image-classification", model=model, image_processor=processor
-        )
+    # 模型的 preprocessor_config.json 中 image_processor_type 为旧版 ViTFeatureExtractor，
+    # AutoImageProcessor 无法识别，直接用具体的 ViT 类加载
+    processor = ViTImageProcessor.from_pretrained(model_name)
+    model = ViTForImageClassification.from_pretrained(model_name)
+    classifier = pipeline(
+        "image-classification", model=model, image_processor=processor
+    )
     return classifier