PyTorch 医学分割

阅读中文简介：点击这里！

❗️🚨 重要通知 🚨❗️

🔗 请参考我们的新仓库：SLab-Medical-Segmentation

🛑 现有仓库已弃用！

备注

我们计划在近期对代码进行重大更新，如果您有任何建议，欢迎随时通过邮件 me 与我联系，或在 issue 中提出。

最新动态

• 2021.1.8 发布了训练和测试代码。
• 2021.2.6 在 Shanshan Li 的帮助下修复了 dice 指标中的一个 bug。
• 2021.2.24 发布了一段视频教程（https://www.bilibili.com/video/BV1gp4y1H7kq/）。
• 2021.5.16 修复了 Unet3D 实现中的一个 bug。
• 2021.5.16 发布了指标代码。
• 2021.6.24 所有参数现在都可以在 hparam.py 中调整。
• 2021.7.7 现在您可以参考 Pytorch-Medical-Classification 中的医学分类任务。
• 2022.5.15 现在您可以参考 SSL-For-Medical-Segmentation 中关于医学分割的半监督学习内容。
• 2022.5.17 我们更新了训练和推理代码，并修复了一些 bug。

环境要求

• pytorch 1.7
• torchio <= 0.18.20
• python >= 3.6

注意事项

• 您可以通过修改 hparam.py 来决定是进行 2D 还是 3D 分割，以及是否支持多类别分割。
• 我们提供了几乎所有 2D 和 3D 分割算法。
• 本仓库兼容几乎所有的医学数据格式（如 nii.gz、nii、mhd、nrrd 等），只需在配置文件的 hparam.py 中修改 fold_arch 即可。建议您在使用前将源图像和标签图像统一转换为同一种格式，并确保标签用 1 标记，而不是 255。
• 如果您希望使用 多类别 分割程序，请自行修改相关代码。我无法识别您的具体类别。
• 无论 2D 还是 3D，本项目都采用 patch 方式处理数据，因此图像不必严格保持相同尺寸。但在 2D 情况下，您应设置足够大的 patch 尺寸。

数据集准备

示例1

如果您的源数据集为：

source_dataset
├── source_1.mhd
├── source_1.zraw
├── source_2.mhd
├── source_2.zraw
├── source_3.mhd
├── source_3.zraw
├── source_4.mhd
├── source_4.zraw
└── ...

而您的标签数据集为：

label_dataset
├── label_1.mhd
├── label_1.zraw
├── label_2.mhd
├── label_2.zraw
├── label_3.mhd
├── label_3.zraw
├── label_4.mhd
├── label_4.zraw
└── ...

那么您需要在 hparam.py 中将 fold_arch 修改为 *.mhd，并将 source_train_dir 设置为 source_dataset，label_train_dir 设置为 label_dataset。

示例2

如果您的源数据集为：

source_dataset
├── 1
    ├── source_1.mhd
    ├── source_1.zraw
├── 2
    ├── source_2.mhd
    ├── source_2.zraw
├── 3
    ├── source_3.mhd
    ├── source_3.zraw
├── 4
    ├── source_4.mhd
    ├── source_4.zraw
└── ...

而您的标签数据集为：

label_dataset
├── 1
    ├── label_1.mhd
    ├── label_1.zraw
├── 2
    ├── label_2.mhd
    ├── label_2.zraw
├── 3
    ├── label_3.mhd
    ├── label_3.zraw
├── 4
    ├── label_4.mhd
    ├── label_4.zraw
└── ...

那么您需要在 hparam.py 中将 fold_arch 修改为 */*.mhd，并将 source_train_dir 设置为 source_dataset，label_train_dir 设置为 label_dataset。

训练

• 不使用预训练模型

将 hparam.train_or_test 设置为 'train'
python main.py

• 使用预训练模型

将 hparam.train_or_test 设置为 'train'
python main.py -k True

推理

• 测试

将 hparam.train_or_test 设置为 'test'
python main.py

示例

教程

• https://www.bilibili.com/video/BV1gp4y1H7kq/

已完成

网络

• 2D

• unet
• unet++
• miniseg
• segnet
• pspnet
• highresnet（复制自 https://github.com/fepegar/highresnet，感谢 fepegar 的慷慨分享！）
• deeplab
• fcn

• 3D

• unet3d
• residual-unet3d
• densevoxelnet3d
• fcn3d
• vnet3d
• highresnert（复制自 https://github.com/fepegar/highresnet，感谢 fepegar 的慷慨分享！）
• densenet3d
• unetr（复制自 https://github.com/tamasino52/UNETR）

指标

• metrics.py 用于评估您的结果

待办事项

• 数据集
• 基准测试
• nnunet

顺便说一句

本项目并不完美，仍存在许多问题。如果您正在使用该项目并希望向作者提供反馈，欢迎发送邮件至 Me。

致谢

本仓库是非官方的 PyTorch 3D 和 2D 医学分割实现，高度参考了 MedicalZooPytorch 和 torchio。感谢上述项目的贡献。本项目在 Prof. Ruoxiu Xiao、Prof. Shuang Song 和 Dr. Cheng Chen 的指导下完成。同时，感谢 Youming Zhang、Daiheng Gao、Jie Zhang、Xing Tao、Weili Jiang 和 Shanshan Li 对我的所有帮助。

Pytorch-Medical-Segmentation 快速上手指南

⚠️ 重要提示：原仓库已停止维护（Deprecated）。作者强烈建议迁移至新仓库：SLab-Medical-Segmentation。以下指南仅针对当前仓库版本供参考。

1. 环境准备

本项目基于 PyTorch 开发，支持多种医学图像格式（如 nii.gz, mhd, nrrd 等）。请确保满足以下系统要求：

• 操作系统: Linux / Windows / macOS
• Python: >= 3.6
• PyTorch: 1.7
• TorchIO: <= 0.18.20

依赖安装命令：

pip install torch==1.7.0
pip install "torchio<=0.18.20"
# 其他常见依赖（如 numpy, scipy 等通常随 torchio 自动安装，若缺失请手动补充）
pip install numpy scipy SimpleITK

💡 国内加速建议：建议使用清华或阿里镜像源加速安装： pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch==1.7.0 "torchio<=0.18.20"

2. 数据准备与配置

在运行代码前，需整理数据集并修改配置文件 hparam.py。

数据集格式要求

• 标签值：请将标签图像中的目标区域标记为 1（不要使用 255）。
• 文件格式：建议将原始图像（Source）和标签图像（Label）转换为相同的格式（如均为 .mhd 或 .nii.gz）。
• 目录结构：支持扁平结构或子目录结构，需在配置中对应调整。

修改配置文件 (hparam.py)

打开 hparam.py，根据实际数据路径和维度进行如下关键设置：

1. 设置训练模式：

   train_or_test = 'train'  # 或 'test'
   

2. 配置数据路径与文件匹配规则：
   • 场景 A（扁平结构）：所有文件直接在文件夹内。

     fold_arch = '*.mhd'  # 根据实际后缀修改，如 '*.nii.gz'
     source_train_dir = 'path/to/source_dataset'
     label_train_dir = 'path/to/label_dataset'
     

   • 场景 B（子目录结构）：每个样本在一个子文件夹内。

     fold_arch = '*/*.mhd' 
     source_train_dir = 'path/to/source_dataset'
     label_train_dir = 'path/to/label_dataset'
     

3. 其他参数：
   • 修改 is_2d 切换 2D/3D 分割。
   • 若进行多分类任务，需自行修改相关代码逻辑以适配类别数量。

3. 基本使用

开始训练

从零开始训练：

python main.py

使用预训练模型微调：

python main.py -k True

推理/测试

完成训练后，将配置项 train_or_test 改为 'test'，然后运行：

# 确保 hparam.py 中 train_or_test = 'test'
python main.py

注意：本项目采用 Patch 机制处理图像，因此输入图像无需严格统一尺寸。但在 2D 模式下，请确保设置的 Patch 尺寸足够大以包含有效特征。