基于 SegFormer 的医学图像语义分割项目
Medical Image Semantic Segmentation Using SegFormer Transformer
本项目使用 NVIDIA 的 SegFormer transformer 模型对医学图像（胃肠道 CT 扫描）进行语义分割，实现对胃部、小肠、大肠三个器官的自动识别和分割。
This project uses NVIDIA's SegFormer transformer model for semantic segmentation of medical images (gastrointestinal CT scans), achieving automatic recognition and segmentation of three organs: stomach, small intestine, and large intestine.
- 🔥 先进架构: 基于 Vision Transformer 的 SegFormer 模型
- ⚡ 高性能训练: 使用 PyTorch Lightning 框架，支持混合精度训练
- 📊 实验追踪: 集成 Weights & Biases 实时监控训练过程
- 🖼️ 数据增强: 使用 Albumentations 进行医学图像增强
- 📈 高精度: 验证集 F1-Score 达到 0.93-0.95
- GPU: NVIDIA GPU with 8GB+ VRAM (推荐 RTX 3060 或更高)
- Python: 3.8+
- CUDA: 11.0+
# 克隆项目
git clone https://github.com/PlutoLei/BMDS_group_project.git
cd BMDS_group_project
# 安装 PyTorch (CUDA 11.8)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
# 安装其他依赖
pip install transformers==4.35.0 lightning==2.1.0 albumentations==1.3.1
pip install torchmetrics==1.2.0 wandb==0.16.0 opencv-python matplotlib# 打开 Jupyter Notebook
jupyter notebook BS6204/Segformer_medical_lightning/Segformer_medical_lightning.ipynb
# 按顺序执行每个单元格
# 数据集会自动下载 (~500MB)BS6204/
└── Segformer_medical_lightning/
├── Segformer_medical_lightning.ipynb # 主训练 Notebook
└── docs/ # 项目文档
├── START_TRAINING.md # 快速开始指南
├── Segformer_医学图像分割_使用说明.md # 完整使用说明
├── Segformer_快速开始.md # 快速入门
├── Segformer_配置推荐表.md # 配置推荐
├── 训练优化改善报告_20251014.md # 训练优化报告
├── 项目优化完整总结_20251014.md # 项目优化总结
└── 项目总结.md # 项目总结
- 深度学习框架: PyTorch 2.0+ & PyTorch Lightning 2.1+
- 模型: SegFormer (NVIDIA) - Vision Transformer
- 数据增强: Albumentations
- 实验追踪: Weights & Biases (WandB)
- 评估指标: TorchMetrics (Dice F1-Score)
- 可视化: Matplotlib, OpenCV
| 指标 Metric | 训练集 Train | 验证集 Validation |
|---|---|---|
| Loss | 0.25-0.30 | 0.30-0.35 |
| F1-Score (Dice) | 0.95-0.97 | 0.93-0.95 |
| 类别 Class | F1-Score | 难度 Difficulty |
|---|---|---|
| 背景 Background | ~0.98 | 简单 Easy |
| 胃部 Stomach | ~0.92-0.94 | 中等 Medium |
| 小肠 Small Intestine | ~0.88-0.92 | 困难 Hard |
| 大肠 Large Intestine | ~0.90-0.93 | 中等 Medium |
详细文档请查看 docs/ 目录：
UW-Madison GI Tract Image Segmentation
- 来源: Kaggle Competition
- 大小: ~500MB
- 类别: 4类 (背景, 胃, 小肠, 大肠)
- 训练集: 3,600+ 图像
- 验证集: 900+ 图像
- 图像尺寸: 288×288 (预处理后)
| 模型 Model | 参数量 Params | 显存需求 VRAM | 推荐用途 Use Case |
|---|---|---|---|
| segformer-b0 | 4M | ~4GB | 快速实验 Quick Test |
| segformer-b1 | 14M | ~6GB | 平衡性能 Balanced |
| segformer-b2 | 28M | ~8GB | 高精度 High Accuracy |
| segformer-b3 | 47M | ~10GB | 生产环境 Production |
| segformer-b4 | 64M | ~12GB | 默认 Default |
| segformer-b5 | 84M | ~16GB | 最高精度 Best |
- 自动从 Kaggle 下载数据集
- 智能数据增强（旋转、翻转、亮度调整等）
- 自动划分训练集和验证集
- 配置化超参数管理
- 自动保存最佳模型
- 实时训练监控和可视化
- 多种评估指标（Dice F1、IoU）
- 可视化分割结果
- 支持批量推理
- WandB 云端实验管理
- TensorBoard 本地可视化
- 训练曲线自动保存
- 医疗用途声明: 本模型仅用于教育和研究目的，不应直接用于临床诊断
- 数据隐私: 使用私有医学数据时请遵守相关法规（HIPAA 等）
- 硬件要求: 建议使用 GPU 训练，CPU 训练速度极慢（50-70小时）
- 版本兼容: Python 3.8-3.10，PyTorch 2.0+
- SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers
- Vision Transformer: An Image is Worth 16x16 Words
- Dice Loss: V-Net: Fully Convolutional Neural Networks
本项目基于 MIT 许可证开源。
引用 Citation:
@article{xie2021segformer,
title={SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers},
author={Xie, Enze and Wang, Wenhai and Yu, Zhiding and Anandkumar, Anima and Alvarez, Jose M and Luo, Ping},
journal={NeurIPS},
year={2021}
}- 原始数据集提供: UW-Madison Carbone Cancer Center
- SegFormer 模型: NVIDIA Research
- 参考教程: LearnOpenCV
⭐ 如果这个项目对您有帮助，请给一个 Star！
If this project helps you, please give it a Star!