fix: 使用 ViT 专用类加载模型,移除 Auto 类兼容回退逻辑
- 将 AutoImageProcessor / AutoModelForImageClassification 替换为
ViTImageProcessor / ViTForImageClassification
- 移除 ValueError 异常捕获回退逻辑,改为直接使用 ViT 类加载,
避免因 preprocessor_config.json 中旧版 image_processor_type
导致 Auto 类无法识别的问题
This commit is contained in:
24
app.py
24
app.py
@@ -12,7 +12,7 @@ from PIL import Image
|
|||||||
# 国内 HuggingFace 镜像加速
|
# 国内 HuggingFace 镜像加速
|
||||||
os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")
|
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 ──────────────────────────────────────────────────
|
# ─── Disease Label Mapping ──────────────────────────────────────────────────
|
||||||
# 模型输出 LABEL_0 ~ LABEL_6,映射为实际病害名称
|
# 模型输出 LABEL_0 ~ LABEL_6,映射为实际病害名称
|
||||||
@@ -91,21 +91,13 @@ html, body, [class*="css"] {
|
|||||||
def load_model():
|
def load_model():
|
||||||
"""加载 HuggingFace 模型(首次运行自动下载,约 343MB)"""
|
"""加载 HuggingFace 模型(首次运行自动下载,约 343MB)"""
|
||||||
model_name = "Dmitry43243242/banana-disease-leaf-model"
|
model_name = "Dmitry43243242/banana-disease-leaf-model"
|
||||||
try:
|
# 模型的 preprocessor_config.json 中 image_processor_type 为旧版 ViTFeatureExtractor,
|
||||||
classifier = pipeline(
|
# AutoImageProcessor 无法识别,直接用具体的 ViT 类加载
|
||||||
"image-classification",
|
processor = ViTImageProcessor.from_pretrained(model_name)
|
||||||
model=model_name,
|
model = ViTForImageClassification.from_pretrained(model_name)
|
||||||
)
|
classifier = pipeline(
|
||||||
except ValueError:
|
"image-classification", model=model, image_processor=processor
|
||||||
# 模型的 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
|
|
||||||
)
|
|
||||||
return classifier
|
return classifier
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user