使用 PyTorch 实现的 3DUNet

简介

该仓库是一个基于 PyTorch 实现的 3DUNet，参考了 此项目。我对代码结构进行了重新设计，并使用该模型在 lits2017 数据集上进行了肝脏及肿瘤分割任务。
论文：3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation

需求：

pytorch >= 1.1.0
torchvision
SimpleITK
Tensorboard
Scipy

代码结构

├── dataset          # 训练与测试数据集
│   ├── dataset_lits_train.py 
│   ├── dataset_lits_val.py
│   ├── dataset_lits_test.py
│   └── transforms.py 
├── models           # 模型设计
│   ├── nn
│   │   └── module.py
│   │── ResUNet.py      # 3DUNet 类
│   │── Unet.py      # 3DUNet 类
│   │── SegNet.py      # 3DUNet 类
│   └── KiUNet.py      # 3DUNet 类
├── experiments           # 训练好的模型
|── utils            # 相关工具
|   ├── common.py
|   ├── weights_init.py
|   ├── logger.py
|   ├── metrics.py
|   └── loss.py
├── preprocess_LiTS.py  # 原始数据预处理
├── test.py          # 测试代码
├── train.py         # 标准训练代码
└── config.py        # 训练与测试的配置信息

快速入门

1) LITS2017 数据集预处理：

1. 从 Google Drive 下载数据集：Liver Tumor Segmentation Challenge.
   或者从我的分享链接下载：https://pan.baidu.com/s/1WgP2Ttxn_CV-yRT4UyqHWw
   提取码：hfl8（数据集分为 batch1 和 batch2 两部分）
2. 解压后将 batch1 和 batch2 合并到一个文件夹中。建议将 LiTS 数据集中 20 个样本（2746）作为测试集，111 个样本（026 和 47~131）作为训练集。请将训练集和测试集的体积数据及分割标签分别放置在不同的本地文件夹中，例如：

raw_dataset:
    ├── test  # 20 个样本（27~46） 
    │   ├── ct
    │   │   ├── volume-27.nii
    │   │   ├── volume-28.nii
    |   |   |—— ...
    │   └── label
    │       ├── segmentation-27.nii
    │       ├── segmentation-28.nii
    |       |—— ...
    │       
    ├── train # 111 个样本（0~26 和 47~131）
    │   ├── ct
    │   │   ├── volume-0.nii
    │   │   ├── volume-1.nii
    |   |   |—— ...
    │   └── label
    │       ├── segmentation-0.nii
    │       ├── segmentation-1.nii
    |       |—— ...

3. 最后，需要在 ./preprocess_LiTS.py 中修改体积数据和分割标签的根路径，例如：

    row_dataset_path = './raw_dataset/train/'  # 原始数据路径
    fixed_dataset_path = './fixed_data/'  # 处理后的数据路径

4. 运行 python ./preprocess_LiTS.py
   如果一切顺利，您将在 ./fixed_data 目录下看到以下文件：

│—— train_path_list.txt
│—— val_path_list.txt
│
|—— ct
│       volume-0.nii
│       volume-1.nii
│       volume-2.nii
│       ...
└─ label
        segmentation-0.nii
        segmentation-1.nii
        segmentation-2.nii
        ...

---

2) 训练 3DUNet

1. 首先，您需要在 config.py 中修改一些参数，尤其是将 --dataset_path 设置为 ./fixed_data。
   文件中的所有参数均已添加注释说明。
2. 其次，运行 python train.py --save model_name
3. 此外，您可以通过浏览器访问 tensorboard --logdir ./output/model_name 来观察训练过程中的 Dice 系数和损失值。

---

3) 测试 3DUNet

运行 test.py
请注意 test.py 中已训练模型的路径。
（由于 3D 卷积运算量较大，我在预测前使用滑动窗口对输入张量进行分块处理，随后将结果拼接以得到最终输出。滑动窗口的大小可在 config.py 中自行设置。）

测试完成后，您可以在对应的文件夹中找到测试结果：./experiments/model_name/result。

您也可以阅读我关于该 3DUNet 项目的中文介绍 这里。不过，我目前不再更新博客，会尽力维护 GitHub 上的代码。
如果您有任何建议或问题，欢迎提交 Issue 与我交流。

3DUNet-Pytorch 快速上手指南

环境准备

• 系统要求：Python 3.6+，Linux/Windows/macOS
• 前置依赖：

  pytorch >= 1.1.0
  torchvision
  SimpleITK
  Tensorboard
  Scipy
  

安装步骤

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple/ torch torchvision SimpleITK tensorboard scipy

基本使用

1. 数据集预处理

• 下载 LITS2017 数据集（百度网盘：https://pan.baidu.com/s/1WgP2Ttxn_CV-yRT4UyqHWw，提取码：hfl8）
• 解压后按以下结构组织数据：

  raw_dataset/
  ├── test/
  │   ├── ct/          # 20个测试样本（27~46）
  │   └── label/
  └── train/
      ├── ct/          # 111个训练样本（0~26, 47~131）
      └── label/
  

• 修改 preprocess_LiTS.py 路径配置：

  row_dataset_path = './raw_dataset/train/'
  fixed_dataset_path = './fixed_data/'
  

• 执行预处理：

  python preprocess_LiTS.py
  

2. 训练模型

• 修改 config.py 中 dataset_path 为 ./fixed_data
• 启动训练：

  python train.py --save liver_segmentation
  

• 监控训练过程：

  tensorboard --logdir ./output/liver_segmentation
  

3. 测试模型

• 直接运行测试脚本（确保 test.py 中模型路径正确）：

  python test.py
  

• 测试结果保存在：./experiments/liver_segmentation/result/