581630c4e0
- 增加模型加载的异常处理，当 pipeline 直接加载失败时回退为
手动加载 AutoImageProcessor 和 AutoModelForImageClassification
- 添加 torchvision 依赖
297 lines
11 KiB
Python
297 lines
11 KiB
Python
"""
|
||
香蕉叶病害智能识别系统
|
||
基于 ViT 视觉变换器的香蕉叶片病害分类应用
|
||
"""
|
||
|
||
import os
|
||
|
||
import streamlit as st
|
||
import plotly.graph_objects as go
|
||
from PIL import Image
|
||
|
||
# 国内 HuggingFace 镜像加速
|
||
os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")
|
||
|
||
from transformers import AutoImageProcessor, AutoModelForImageClassification, pipeline # noqa: E402
|
||
|
||
# ─── Disease Label Mapping ──────────────────────────────────────────────────
|
||
# 模型输出 LABEL_0 ~ LABEL_6，映射为实际病害名称
|
||
DISEASE_INFO = {
|
||
"LABEL_0": {
|
||
"name": "细菌性枯萎病",
|
||
"name_en": "Banana Bacterial Wilt (BBW)",
|
||
"color": "#e74c3c",
|
||
"description": "由细菌引起，叶片从边缘开始变黄枯萎，最终整株死亡。",
|
||
"advice": "及时清除病株，避免刀具交叉感染，使用无菌工具操作，加强田间排水。",
|
||
},
|
||
"LABEL_1": {
|
||
"name": "黑条斑病",
|
||
"name_en": "Black Sigatoka",
|
||
"color": "#8e44ad",
|
||
"description": "真菌性病害，叶片出现黑色条纹，严重时叶片大面积枯死。",
|
||
"advice": "选用抗病品种，合理密植保持通风，及时清除病叶，适时喷施杀菌剂。",
|
||
},
|
||
"LABEL_2": {
|
||
"name": "苞片花叶病毒病",
|
||
"name_en": "Bract Mosaic Virus (BBMV)",
|
||
"color": "#e67e22",
|
||
"description": "病毒性病害，苞片和叶片出现花叶状斑纹，影响果实品质。",
|
||
"advice": "使用无病种苗，防治传毒蚜虫，及时拔除病株，避免机械传播。",
|
||
},
|
||
"LABEL_3": {
|
||
"name": "健康叶片",
|
||
"name_en": "Healthy",
|
||
"color": "#27ae60",
|
||
"description": "叶片状态良好，无明显病害症状，颜色鲜绿。",
|
||
"advice": "继续保持良好的田间管理，合理施肥灌溉，定期巡检。",
|
||
},
|
||
"LABEL_4": {
|
||
"name": "花叶病毒病",
|
||
"name_en": "Banana Mosaic Virus",
|
||
"color": "#f39c12",
|
||
"description": "病毒性病害，叶片出现黄绿相间的花叶症状，植株矮化。",
|
||
"advice": "选用无毒组培苗，控制蚜虫等传毒媒介，拔除病株，工具消毒。",
|
||
},
|
||
"LABEL_5": {
|
||
"name": "巴拿马病（枯萎病）",
|
||
"name_en": "Panama Disease (Fusarium Wilt)",
|
||
"color": "#c0392b",
|
||
"description": "由尖孢镰刀菌引起，维管束变褐，叶片从下部开始黄化枯萎。",
|
||
"advice": "种植抗病品种，严格检疫，病区轮作或休耕，避免土壤传播。",
|
||
},
|
||
"LABEL_6": {
|
||
"name": "黄条斑病",
|
||
"name_en": "Yellow Sigatoka",
|
||
"color": "#d4ac0d",
|
||
"description": "真菌性病害，叶片出现黄色条纹，比黑条斑病症状较轻。",
|
||
"advice": "保持田间通风透光，清除病残叶，适时用药防治，合理施肥增强抗性。",
|
||
},
|
||
}
|
||
|
||
|
||
# ─── Page Config ────────────────────────────────────────────────────────────
|
||
st.set_page_config(
|
||
page_title="香蕉叶病害识别",
|
||
page_icon="🍌",
|
||
layout="wide",
|
||
initial_sidebar_state="expanded",
|
||
)
|
||
|
||
st.markdown("""
|
||
<style>
|
||
html, body, [class*="css"] {
|
||
font-family: "PingFang SC", "Microsoft YaHei", "Noto Sans SC", sans-serif;
|
||
}
|
||
</style>
|
||
""", unsafe_allow_html=True)
|
||
|
||
|
||
# ─── Model Loading ──────────────────────────────────────────────────────────
|
||
@st.cache_resource
|
||
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
|
||
)
|
||
return classifier
|
||
|
||
|
||
# ─── Sidebar ────────────────────────────────────────────────────────────────
|
||
with st.sidebar:
|
||
st.header("🍌 香蕉叶病害识别")
|
||
st.caption("上传叶片照片，AI 帮你识别病害")
|
||
st.divider()
|
||
|
||
st.subheader("📷 上传图片")
|
||
uploaded_file = st.file_uploader(
|
||
"选择香蕉叶片图片",
|
||
type=["jpg", "jpeg", "png", "webp"],
|
||
help="支持 JPG、PNG、WebP 格式",
|
||
)
|
||
|
||
st.divider()
|
||
st.subheader("ℹ️ 使用说明")
|
||
st.markdown("""
|
||
1. 上传一张香蕉叶片照片
|
||
2. 系统自动识别病害类型
|
||
3. 查看各类别置信度
|
||
4. 获取防治建议
|
||
|
||
**Tips:**
|
||
- 图片越清晰，识别越准确
|
||
- 尽量拍摄完整叶片
|
||
- 病变部位清晰可见
|
||
""")
|
||
|
||
st.divider()
|
||
st.subheader("📋 可识别病害")
|
||
for label_key, info in DISEASE_INFO.items():
|
||
st.markdown(
|
||
f"<span style='color:{info['color']}'>●</span> **{info['name']}** ({info['name_en']})",
|
||
unsafe_allow_html=True,
|
||
)
|
||
|
||
|
||
# ─── Main Content ───────────────────────────────────────────────────────────
|
||
st.title("🍌 香蕉叶病害智能识别系统")
|
||
st.caption("基于 ViT 视觉变换器 · 7 类香蕉叶病害自动分类")
|
||
|
||
# 加载模型
|
||
with st.spinner("正在加载模型（首次运行需下载约 343MB）..."):
|
||
classifier = load_model()
|
||
|
||
if uploaded_file is not None:
|
||
# 显示上传的图片
|
||
image = Image.open(uploaded_file).convert("RGB")
|
||
|
||
col_img, col_result = st.columns([1, 1])
|
||
|
||
with col_img:
|
||
st.subheader("🖼️ 上传图片")
|
||
st.image(image, use_container_width=True)
|
||
|
||
# 推理
|
||
with st.spinner("正在识别病害..."):
|
||
results = classifier(image)
|
||
|
||
# 解析结果
|
||
top_result = results[0]
|
||
top_label = top_result["label"]
|
||
top_score = top_result["score"]
|
||
disease = DISEASE_INFO.get(top_label, {
|
||
"name": top_label,
|
||
"name_en": top_label,
|
||
"color": "#95a5a6",
|
||
"description": "未知病害类型",
|
||
"advice": "建议咨询专业植保人员",
|
||
})
|
||
|
||
with col_result:
|
||
st.subheader("🔍 识别结果")
|
||
is_healthy = top_label == "LABEL_3"
|
||
if is_healthy:
|
||
st.success(f"**{disease['name']}** — 置信度 {top_score:.1%}")
|
||
else:
|
||
st.warning(f"**{disease['name']}** — 置信度 {top_score:.1%}")
|
||
|
||
st.markdown(f"**英文名**: {disease['name_en']}")
|
||
st.markdown(f"**症状描述**: {disease['description']}")
|
||
|
||
# ─── Confidence Bar Chart ────────────────────────────────────────────────
|
||
st.subheader("📊 各类别置信度")
|
||
|
||
labels = [r["label"] for r in results]
|
||
scores = [r["score"] for r in results]
|
||
display_names = [
|
||
f"{DISEASE_INFO.get(l, {}).get('name', l)}" for l in labels
|
||
]
|
||
colors = [
|
||
DISEASE_INFO.get(l, {}).get("color", "#95a5a6") for l in labels
|
||
]
|
||
|
||
fig = go.Figure()
|
||
fig.add_trace(go.Bar(
|
||
x=scores,
|
||
y=display_names,
|
||
orientation="h",
|
||
marker=dict(
|
||
color=colors,
|
||
opacity=0.85,
|
||
line=dict(color="rgba(0,0,0,0.08)", width=1),
|
||
),
|
||
text=[f"{s:.1%}" for s in scores],
|
||
textposition="outside",
|
||
textfont=dict(color="#5a5a5a", size=11),
|
||
))
|
||
fig.update_layout(
|
||
xaxis=dict(
|
||
title="置信度",
|
||
range=[0, 1.15],
|
||
tickformat=".0%",
|
||
color="#5a5a5a",
|
||
gridcolor="rgba(0,0,0,0.06)",
|
||
),
|
||
yaxis=dict(color="#2c2c2c", autorange="reversed"),
|
||
paper_bgcolor="rgba(0,0,0,0)",
|
||
plot_bgcolor="rgba(0,0,0,0)",
|
||
font=dict(color="#2c2c2c", size=11),
|
||
margin=dict(t=10, b=30, l=130, r=60),
|
||
height=320,
|
||
showlegend=False,
|
||
)
|
||
st.plotly_chart(fig, use_container_width=True)
|
||
|
||
# ─── Pie Chart ──────────────────────────────────────────────────────────
|
||
col_pie, col_advice = st.columns([1, 1])
|
||
|
||
with col_pie:
|
||
st.subheader("📈 置信度分布")
|
||
fig_pie = go.Figure()
|
||
fig_pie.add_trace(go.Pie(
|
||
labels=display_names,
|
||
values=scores,
|
||
marker=dict(colors=colors),
|
||
textinfo="label+percent",
|
||
textfont=dict(size=10),
|
||
hole=0.45,
|
||
))
|
||
fig_pie.update_layout(
|
||
paper_bgcolor="rgba(0,0,0,0)",
|
||
font=dict(color="#2c2c2c", size=10),
|
||
margin=dict(t=10, b=10, l=10, r=10),
|
||
height=300,
|
||
showlegend=False,
|
||
)
|
||
st.plotly_chart(fig_pie, use_container_width=True)
|
||
|
||
# ─── Disease Advice ─────────────────────────────────────────────────────
|
||
with col_advice:
|
||
st.subheader("💡 防治建议")
|
||
if is_healthy:
|
||
st.success(f"叶片状态健康，继续保持良好的田间管理！")
|
||
else:
|
||
st.error(f"检测到 **{disease['name']}** 病害")
|
||
st.info(disease["advice"])
|
||
|
||
# 如果置信度不高，给出提示
|
||
if top_score < 0.6:
|
||
st.warning(
|
||
"⚠️ 置信度较低，建议：\n"
|
||
"- 确保图片清晰且病变部位可见\n"
|
||
"- 尝试拍摄完整叶片\n"
|
||
"- 可咨询专业植保人员确认"
|
||
)
|
||
|
||
else:
|
||
# 无图片上传时的占位内容
|
||
st.info("👈 请在左侧上传一张香蕉叶片图片开始识别")
|
||
|
||
# 展示示例区域
|
||
st.subheader("📖 支持识别的病害类型")
|
||
cols = st.columns(4)
|
||
for i, (label_key, info) in enumerate(DISEASE_INFO.items()):
|
||
with cols[i % 4]:
|
||
with st.container(border=True):
|
||
st.markdown(
|
||
f"<div style='color:{info['color']}; font-size:24px; font-weight:bold;'>●</div>",
|
||
unsafe_allow_html=True,
|
||
)
|
||
st.markdown(f"**{info['name']}**")
|
||
st.caption(info["name_en"])
|
||
|
||
st.divider()
|
||
st.caption("香蕉叶病害智能识别系统 · 基于 ViT (Vision Transformer) · 仅供参考，确诊请咨询专业植保人员")
|