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

View File

@@ -10,6 +10,7 @@ dependencies = [
"ruff>=0.14.8", "ruff>=0.14.8",
"streamlit==1.52.1", "streamlit==1.52.1",
"torch>=2.7.0", "torch>=2.7.0",
"torchvision>=0.22.0",
"transformers>=4.52.0", "transformers>=4.52.0",
"pillow>=11.2.0", "pillow>=11.2.0",
] ]

22
uv.lock generated
View File

@@ -59,6 +59,7 @@ dependencies = [
{ name = "ruff" }, { name = "ruff" },
{ name = "streamlit" }, { name = "streamlit" },
{ name = "torch" }, { name = "torch" },
{ name = "torchvision" },
{ name = "transformers" }, { name = "transformers" },
] ]
@@ -70,6 +71,7 @@ requires-dist = [
{ name = "ruff", specifier = ">=0.14.8" }, { name = "ruff", specifier = ">=0.14.8" },
{ name = "streamlit", specifier = "==1.52.1" }, { name = "streamlit", specifier = "==1.52.1" },
{ name = "torch", specifier = ">=2.7.0" }, { name = "torch", specifier = ">=2.7.0" },
{ name = "torchvision", specifier = ">=0.22.0" },
{ name = "transformers", specifier = ">=4.52.0" }, { name = "transformers", specifier = ">=4.52.0" },
] ]
@@ -1159,6 +1161,26 @@ wheels = [
{ url = "https://mirrors.aliyun.com/pypi/packages/cf/bf/c8d12a2c86dbfd7f40fb2f56fbf5a505ccf2d9ce131eb559dfc7c51e1a04/torch-2.11.0-cp314-cp314t-win_amd64.whl", hash = "sha256:b2a43985ff5ef6ddd923bbcf99943e5f58059805787c5c9a2622bf05ca2965b0" }, { url = "https://mirrors.aliyun.com/pypi/packages/cf/bf/c8d12a2c86dbfd7f40fb2f56fbf5a505ccf2d9ce131eb559dfc7c51e1a04/torch-2.11.0-cp314-cp314t-win_amd64.whl", hash = "sha256:b2a43985ff5ef6ddd923bbcf99943e5f58059805787c5c9a2622bf05ca2965b0" },
] ]
[[package]]
name = "torchvision"
version = "0.26.0"
source = { registry = "https://mirrors.aliyun.com/pypi/simple" }
dependencies = [
{ name = "numpy" },
{ name = "pillow" },
{ name = "torch" },
]
wheels = [
{ url = "https://mirrors.aliyun.com/pypi/packages/7f/c8/9bffa9c7f7bdf95b2a0a2dc535c290b9f1cc580c3fb3033ab1246ffffdeb/torchvision-0.26.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:eb61804eb9dbe88c5a2a6c4da8dec1d80d2d0a6f18c999c524e32266cb1ebcd3" },
{ url = "https://mirrors.aliyun.com/pypi/packages/7b/ac/48f28ffd227991f2e14f4392dde7e8dc14352bb9428c1ef4a4bbf5f7ed85/torchvision-0.26.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:9a904f2131cbfadab4df828088a9f66291ad33f49ff853872aed1f86848ef776" },
{ url = "https://mirrors.aliyun.com/pypi/packages/a4/21/a2266f7f1b0e58e624ff15fd6f01041f59182c49551ece0db9a183071329/torchvision-0.26.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:0f3e572efe62ad645017ea847e0b5e4f2f638d4e39f05bc011d1eb9ac68d4806" },
{ url = "https://mirrors.aliyun.com/pypi/packages/fc/ba/1666f90bc0bdd77aaa11dcc42bb9f621a9c3668819c32430452e3d404730/torchvision-0.26.0-cp314-cp314-win_amd64.whl", hash = "sha256:114bec0c0e98aa4ba446f63e2fe7a2cbca37b39ac933987ee4804f65de121800" },
{ url = "https://mirrors.aliyun.com/pypi/packages/45/8f/1f0402ac55c2ae15651ff831957d083fe70b2d12282e72612a30ba601512/torchvision-0.26.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:b7d3e295624a28b3b1769228ce1345d94cf4d390dd31136766f76f2d20f718da" },
{ url = "https://mirrors.aliyun.com/pypi/packages/d2/6a/18a582fe3c5ee26f49b5c9fb21ad8016b4d1c06d10178894a58653946fda/torchvision-0.26.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:7058c5878262937e876f20c25867b33724586aa4499e2853b2d52b99a5e51953" },
{ url = "https://mirrors.aliyun.com/pypi/packages/c5/9b/f7e119b59499edc00c55c03adc9ec3bd96144d9b81c46852c431f9c64a9a/torchvision-0.26.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:8008474855623c6ba52876589dc52df0aa66e518c25eca841445348e5f79844c" },
{ url = "https://mirrors.aliyun.com/pypi/packages/d0/6a/09f3844c10643f6c0de5d95abc863420cfaf194c88c7dffd0ac523e2015f/torchvision-0.26.0-cp314-cp314t-win_amd64.whl", hash = "sha256:e9d0e022c19a78552fb055d0414d47fecb4a649309b9968573daea160ba6869c" },
]
[[package]] [[package]]
name = "tornado" name = "tornado"
version = "6.5.5" version = "6.5.5"