fix: 兼容模型加载失败场景,手动构造 image-classification pipeline
- 增加模型加载的异常处理,当 pipeline 直接加载失败时回退为
手动加载 AutoImageProcessor 和 AutoModelForImageClassification
- 添加 torchvision 依赖
This commit is contained in:
22
app.py
22
app.py
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user