fix: 兼容模型加载失败场景,手动构造 image-classification pipeline

- 增加模型加载的异常处理,当 pipeline 直接加载失败时回退为
    手动加载 AutoImageProcessor 和 AutoModelForImageClassification
  - 添加 torchvision 依赖
This commit is contained in:
zhenghu
2026-04-15 14:27:58 +08:00
parent 86541eb55e
commit 581630c4e0
3 changed files with 40 additions and 5 deletions

22
app.py
View File

@@ -12,7 +12,7 @@ from PIL import Image
# 国内 HuggingFace 镜像加速
os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")
from transformers import pipeline # noqa: E402
from transformers import AutoImageProcessor, AutoModelForImageClassification, pipeline # noqa: E402
# ─── Disease Label Mapping ──────────────────────────────────────────────────
# 模型输出 LABEL_0 ~ LABEL_6映射为实际病害名称
@@ -90,10 +90,22 @@ html, body, [class*="css"] {
@st.cache_resource
def load_model():
"""加载 HuggingFace 模型(首次运行自动下载,约 343MB"""
classifier = pipeline(
"image-classification",
model="Dmitry43243242/banana-disease-leaf-model",
)
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
)
return classifier