基于EfficientNetV2的中草药图像识别系统开发

1. 项目概述与背景

中草药识别一直是中医药数字化进程中的重要课题。传统的人工鉴别方法依赖药师经验，存在主观性强、效率低下等问题。我们基于PyTorch框架和EfficientNetV2网络，构建了一个能够自动识别5种常见中草药（百合、党参、山魈、枸杞、槐花、金银花）的深度学习系统。这个项目不仅实现了90%以上的分类准确率，还开发了用户友好的GUI界面，为中医药信息化提供了实用工具。

在医疗AI领域，图像分类技术的应用越来越广泛。相比传统CNN网络，EfficientNetV2在保持高精度的同时大幅提升了训练和推理速度，这对资源受限的中小型医疗机构特别有价值。我们的实现方案在消费级GPU上即可运行，单张图片的推理时间控制在200ms以内。

2. 数据准备与预处理

2.1 数据集构建

我们收集了6类共900张中草药高清图片，每类约150张。数据采集时特别注意了以下几点：

• 多角度拍摄：包含药材的正面、侧面、横截面等视角
• 多光照条件：自然光、室内光、闪光灯等不同光源
• 多背景环境：纯色背景和实际使用场景

数据集目录结构如下：

code复制ChineseMedicine/
├── 百合/
│   ├── image_001.jpg
│   └── ...
├── 党参/
├── 山魈/ 
├── 枸杞/
├── 槐花/
└── 金银花/

2.2 数据增强策略

为提高模型泛化能力，我们设计了以下数据增强方案：

python复制data_transform = {
    "train": transforms.Compose([
        transforms.RandomResizedCrop(300),  # 随机裁剪缩放
        transforms.RandomHorizontalFlip(),  # 水平翻转
        transforms.ColorJitter(brightness=0.2, contrast=0.2), # 颜色扰动
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ]),
    "val": transforms.Compose([
        transforms.Resize(384),
        transforms.CenterCrop(384),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ])
}

注意事项：中草药图像增强需要保持药材的关键特征不被破坏。例如，枸杞的纹理、金银花的形状等关键特征应在增强后仍然清晰可辨。

3. EfficientNetV2模型详解

3.1 网络架构创新

EfficientNetV2相比V1版本主要做了以下改进：

1. Fused-MBConv模块：将V1中的深度可分离卷积替换为常规卷积+BN+Swish的组合，在浅层网络中效率更高
2. 渐进式训练策略：训练初期使用小尺寸图像，后期逐步增大
3. 自适应正则化：根据图像尺寸动态调整Dropout和RandAugment强度

网络结构对比如下：

3.2 PyTorch实现要点

我们基于官方预训练模型进行微调，关键代码如下：

python复制from torchvision.models import efficientnet_v2_s

def create_model(num_classes=6):
    model = efficientnet_v2_s(pretrained=True)
    # 修改最后一层全连接
    model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
    return model

实操技巧：冻结除分类头外的所有层可以显著加快训练速度。当验证集准确率停滞时再解冻部分深层网络。

4. 模型训练与优化

4.1 训练参数配置

我们采用以下超参数设置：

• 初始学习率：0.01（使用余弦退火衰减）
• 批量大小：4（受GPU显存限制）
• 训练轮次：100
• 优化器：SGD（动量0.9，权重衰减1e-4）

学习率调度策略采用余弦退火：

python复制lf = lambda x: ((1 + math.cos(x * math.pi / args.epochs)) / 2) * (1 - args.lrf) + args.lrf
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)

4.2 训练过程监控

使用TensorBoard记录训练指标：

python复制tb_writer.add_scalar("train_loss", train_loss, epoch)
tb_writer.add_scalar("val_acc", val_acc, epoch) 

典型训练曲线显示：

• 约30轮后验证准确率趋于稳定
• 最佳模型出现在第65轮左右（验证准确率92.3%）
• 没有出现过拟合现象（训练/验证损失同步下降）

5. 系统集成与GUI开发

5.1 模型部署方案

将训练好的PyTorch模型转换为TorchScript格式，便于生产环境部署：

python复制model = create_model()
model.load_state_dict(torch.load('best_model.pth'))
model.eval()
traced_script = torch.jit.script(model)
traced_script.save('medicine_classifier.pt')

5.2 PyQt5界面设计

GUI主要功能模块：

• 图片加载区域：支持拖拽和文件选择
• 结果显示面板：显示分类结果和置信度
• 历史记录查询：保存最近10次识别记录

核心预测代码：

python复制def predict_image(image_path):
    img = Image.open(image_path)
    img_tensor = val_transform(img).unsqueeze(0)
    with torch.no_grad():
        output = model(img_tensor)
    prob = torch.nn.functional.softmax(output, dim=1)
    pred_idx = torch.argmax(prob).item()
    return class_names[pred_idx], prob[0][pred_idx].item()

6. 性能优化与实际问题解决

6.1 常见问题排查

1. CUDA内存不足：

   • 减小批量大小（可降至2）
   • 使用混合精度训练

   python复制scaler = torch.cuda.amp.GradScaler()
   with torch.cuda.amp.autocast():
       outputs = model(inputs)
   

2. 类别不平衡：

   • 采用加权交叉熵损失
   • 对少数类进行过采样

3. 推理速度慢：

   • 使用ONNX Runtime加速
   • 量化模型（FP16或INT8）

6.2 效果提升技巧

• 关键区域增强：对药材特定部位（如枸杞的果柄）进行局部放大
• 多模型集成：结合ResNet和MobileNet的结果进行投票
• 测试时增强(TTA)：对输入图像进行多种变换后取平均结果

7. 项目扩展方向

在实际部署中，我们发现几个有价值的改进点：

1. 移动端适配：将模型转换为TFLite格式，开发Android/iOS应用
2. 细粒度分类：区分药材的产地、等级等更细维度
3. 异常检测：识别霉变、虫蛀等质量问题
4. 3D识别：结合多视角图像提升准确率

这个项目完整代码已开源，包含训练脚本、预训练模型和GUI实现。对于想入门医疗AI的开发者，建议先从少量药材类别开始，逐步扩展分类体系。在实际应用中，持续收集真实场景数据对模型迭代至关重要。