nnUNet图像分割教程！从环境配置到训练与推理-轻识

向AI转型的程序员都关注了这个号👇👇👇

nnUNet是什么？

nnU-Net是由德国癌症研究中心、海德堡大学以及海德堡大学医院研究人员(Fabian Isensee, Jens Petersen, Andre Klein)提出来的一个自适应任何新数据集的医学影像分割框架，该框架能根据给定数据集的属性自动调整所有超参数，整个过程无需人工干预。仅仅依赖于朴素的U-Net结构（就是原始U-Net）和鲁棒的训练方案，nnU-Net在六个得到公认的分割挑战中实现了最先进的性能。

关于nnUNet，如果你想更加深入的了解，我推荐看一下这论文

https://arxiv.org/abs/1809.10486

更深入的学习还是建议精读论文。

一、配置虚拟环境

首先，我们想实现nnUNet的使用，需要配置类似如下的环境：

上图所示为nnUNet的官方readme文档建议的环境配置，如果你对cuda，cudnn，torch之类的词语还不算熟悉，可以看一下我的另一篇博客用人话讲解深度学习中CUDA，cudatookit，cudnn和pytorch的关系

https://blog.csdn.net/m0_68239345/article/details/128803487?spm=1001.2014.3001.5501

至于为什么要在虚拟环境中安装nnUNet，当我们在安装包的时候，pip install django安装一个包会附带安装数个其他关联的包。但是当我们删除这个包的时候，我们之后仅仅删除这一个包，安装时附带的其他包并不会删除。**简单来说，当我们的环境使用时间长的时候，我们的包是很难进行管理的。**所以非常建议在一个全新的虚拟环境中来完成接下来的操作！！！

二、安装nnUNet框架

此时，你应该已经配置好了自己的虚拟环境并且可以打印出类似上图的环境版本（可以更高，但最好不要太低），注意，后续的操作均要在你激活你想要使用的那个虚拟环境的前提下进行！！！！

1.安装nnUNet

根据readme文档，这里应该有两种方案可供选择

（1）用作标准化基线、开箱即用的分割算法或使用预训练模型进行推理:

pip install nnunet

（2）用作集成框架（这将在您的计算机上创建nnU-Net代码的副本，以便您可以根据需要对其进行修改）

你想把nnUNet的文件夹放在哪，就在哪个路径下运行这些命令！

git clone https://github.com/MIC-DKFZ/nnUNet.git

cd nnUNet

pip install -e .#最后这个点也不能忽略

对于我来说，因为我后续是要改网络代码的，所以我选择第二种方法，下面详细说明一下这三行命令都是什么意思：

git clone其实就是把人家github上的代码克隆过来，这一步其实和直接复制粘贴代码文件是一样的，总共也就1Mb

cd nnUNet不用说了，就是进入文件夹

其实这个时候就已经有一个nnUNet完整的文件夹了，里面包含这些东西：

最后pip install -e .相当于python setup.py,也就是运行上图这个setup.py文件

这个文件是用来干什么的呢？

安装nnUNet需要的python包

向终端添加几个新命令。这些命令用于运行整个nnU-Net pipeline。您可以从系统上的任何位置执行它们。所有nnU-Net命令都带有前缀“nnUNet_”，以便于识别。

这一步我遇到的两个问题：

（1）从github上git clone代码的时候速度慢的离谱（20kb/s），这个问题主要是github的域名在国内被限制了，网上有很多方法解决，比如下面这个网址可以参考https://www.jianshu.com/p/d58ab49ba98b

（2）运行pip命令的时候速度很慢，这个问题导致我运行pip install -e .的时候超时报错了，同样给出解决方法的链接解决Linux，Ubuntu下使用python包管理工具pip命令安装和下载包速度很慢、失败或者connection timeout等问题_一点儿也不萌的萌萌的博客-CSDN博客_linux 下载python failed: connection timed out.

2.安装隐藏层hiddenlayer（可选）

隐藏层使nnU-net能够给出其生成的网络拓扑图（后面会细说），安装命令如下（这是一整行命令，请务必一起复制粘贴）

pip install

--upgrade git+https://github.com/FabianIsensee/hiddenlayer.git@more_plotted_details#egg=hiddenlayer

这里我遇到的唯一问题就是上面说的pip命令速度太慢，也是根据上述解决方案来解决的。

三、数据集准备

nnUNet对于你要训练的数据是有严格要求的，这第一点就体现在我们保存数据的路径上，请初学者务必按照我下面的样式来创建相应的文件夹并存入数据！！！

第一步：你现在应该有一个名为nnUNet的文件夹（上面有图），进入它，在里面创建一个名为nnUNetFrame的文件夹

第二步：在nnUNetFrame文件夹中创建一个名为DATASET的文件夹，后面我们会用它来存放数据

第三步：在DATASET文件夹中创建三个文件夹，它们分别是nnUNet_raw，nnUNet_preprocessed，nnUNet_trained_models

第四步：进入上面第二个文件夹nnUNet_raw，创建nnUNet_cropped_data文件夹和nnUNet_raw_data文件夹，右边存放原始数据，左边存放crop以后的数据。

第五步：进入右边文件夹nnUNet_raw_data，创建一个名为Task01_BrainTumour的文件夹（解释：这个Task01_BrainTumour是nnUNet的作者参加的一个十项全能竞赛的子任务名，也是我要实践的分割任务，类似的还有Task02_Heart，就是分割心脏的。如果你想分割自己的数据集，建议Task_id从500开始，这样以确保不会与nnUNet的预训练模型发生冲突（ID不能超过999））

第六步：将下载好的公开数据集或者自己的数据集放在上面创建好的任务文件夹下，下面还以Task01_BrainTumour竞赛为例，解释下数据应该怎么存放和编辑：

进入这个网站http://medicaldecathlon.com/.下载对应的数据集（<–网上学科议建<–），取代上面你自己创建的Task01_BrainTumour文件夹。

你会发现目录是这个样子的：json文件是对三个文件夹内容的字典呈现（关乎你的训练），imagesTr是你的训练数据集，打开后你会发现很多的有序的nii.gz的训练文件，而labelsTr里时对应这个imagesTr的标签文件，同样为nii.gz。目前只能是nii.gz文件，nii文件都不行。训练阶段的imageTs文件夹先不管，其实这个文件夹出现在任何位置都可以。（解释：nnUNet使用的是五折交叉验证，并没有验证集）

四、设置nnUNet读取文件的路径

nnUNet是如何知道你的文件存放在哪儿呢，当然要在环境中创建一个路径，这个路径你唯一需要更改的是/nnUNet之前的路径，因为后面的路径你和我是一样的。

第一步：在home目录下按ctrl + h，显示隐藏文件

第二步：找到.bashrc文件，打开

第三步：在文档末尾添加下面三行，保存文件。

第四步：在home下打开终端，输入source .bashrc来更新该文档

现在nnUNet已经知道怎么读取你的文件了。

五、数据集转换

1.数据集转换是什么，为什么要进行数据集转换？

nnUNet要求将原始数据转换成特定的格式，以便了解如何读取和解释数据。

每个分割数据集存储为单独的“任务”。命名包括任务与任务ID，即三位整数和相关联的任务名称。

比如Task001_BrainTumour的任务名称为“脑瘤”，任务ID为1。

在每个任务文件夹中，预期的结构如下：

图像可能具有多种模态，这对于医学图像来说尤其常见。

nnU-Net通过其后缀（文件名末尾的四位整数）识别成像模态。因此，图像文件必须遵循以下命名约定：case_identifier_XXXX.nii.gz

这里，XXXX是模态标识符。dataset.json文件中指定了这些标识符所属的模态

标签文件保存为case_identifier.nii.gz

此命名方案产生以下文件夹结构。用户有责任将其数据转换为这种格式！

下面是MSD的第一个任务的示例：BrainTumor。每个图像有四种模态：FLAIR（0000）、T1w（0001）、T1gd（0002）和T2w（0003）。请注意，imagesTs文件夹是可选的，不必存在

如果对于数据集转换这件事还是不明白，打开你的nnUNet文件夹，在/nnUNet/documentation/文件夹下找到dataset_conversion这个文件来进一步学习。

2.运行数据集转换的命令

依旧以Task01_BrainTumour为例：

nnUNet_convert_decathlon_task -i /home/work/nnUNet/nnUNetFrame/DATASET/nnUNet_raw/nnUNet_raw_data/Task01_BrainTumour

转换操作完成以后，你会发现在你的Task01_BrainTumour文件夹旁边，出现了一个Task001_BrainTumour文件夹，打开看一下，里面的格式应该和我上面展示的一样

3.关于dataset.json文件

这个文件包含你的训练数据信息和任务信息，如果你按照我的建议下载了Task01的数据集，那里面是包含dataset.json文件的，如果你有训练自己的数据集的需求，在我的另一篇博客里会有详细的说明。

六、数据预处理

nnUNet_plan_and_preprocess -t 1 --verify_dataset_integrity

只需要一行命令，因为我们的Task_id是1，所以这里的数字就是1。这个过程会消耗很多的时间，速度慢的原因在于对要进行插值等各种操作。

根据nnUNet框架，三维医学图像分割的通用预处理可以分为四步，分别是数据格式的转换，裁剪crop，重采样resample以及标准化normalization。如果你想进一步学习，推荐学习这篇文章如何针对三维医学图像分割任务进行通用数据预处理：nnUNet中预处理流程总结及代码分析

https://zhuanlan.zhihu.com/p/335970913

运行“nnUNet_plan_and_preprocess”将使用预处理数据填充文件夹

我们将在nnUNet_preprocessed/Task001_BrainTumour中找到这条命令的输出结果。使用2D U-Net以及所有适用的3D U-Net的预处理数据创建子文件夹。它还将为2D和3D配置创建“plans”文件（结尾为.pkl）。这些文件包含生成的分割 pipeline 配置，将由nnUNetTrainer读取（见下文）。请注意，预处理的数据文件夹仅包含训练案例。测试图像没有经过预处理。测试集的预处理将会在推理过程中实时进行。

另外，`–verify_dataset_integrity”应至少在给定数据集上首次运行命令时运行。这将对数据集执行一些检查，以确保其与nnU-Net兼容。如果此检查通过一次，则可以在以后的运行中省略。如果您遵守数据集转换指南（请参见上文），那么这条命令一定会通过的。

七、模型训练

1.写在训练前：更改epoch

nnUNet默认原始的epoch是1000，这太久了，我们打开nnUNet/nnunet/training/network_training/nnUNetTrainerV2.p

第48行的max_epoch来修改epoch

2.关于训练的运行命令

nnU-Net在5倍交叉验证中训练所有U-Net配置。这使nnU-Net能够确定训练数据集的后处理和集合（参见下文）

我们在数据预处理那一步的时候创建了三个U-Net配置：2D U-Net、对全分辨率图像进行操作的3D U-Net以及3D U-Net级联，其中级联的第一个U-Net在下采样图像中创建粗分割图，然后由第二个U-Net进行细化。我们在训练的时候可以自由选用它们。

训练模型使用“nnUNet_train”命令完成。命令的一般结构为：

nnUNet_train CONFIGURATION TRAINER_CLASS_NAME TASK_NAME_OR_ID FOLD --npz (additional options)

CONFIGURATION是一个字符串，用于标识所请求的U-Net配置。

TRAINER_CLASS_NAME是model trainer的名称。如果您实施定制trainers（nnU-Net作为一个框架），您可以在此处指定您的定制trainers。

TASK_NAME_OR_ID指定应训练的数据集，FOLD指定训练的是5倍交叉验证的哪一倍。

“–npz”使模型在最终验证期间保存softmax输出。它仅适用于计划在之后运行“nnUNet_find_best_configuration”的训练

（这是nnU Nets自动选择最佳性能（集合）配置，见下文）。

对于我们的Task01来说，应该运行的命令如下

nnUNet_train 3d_fullres nnUNetTrainerV2 1 0 --npz

3d_fullres代表我们选用对全分辨率图像进行操作的3D U-Net

nnUNetTrainerV2是我们选用的训练器

1代表你的任务ID

0代表五折交叉验证中的第0折

下面给出各种配置的nnUNet网络需要的训练命令

2D U-Net

For FOLD in [0, 1, 2, 3, 4], run:

注意，级联的3D全分辨率U-Net需要预先完成低分辨率U-Net的five folds！

3.训练结果

训练后的模型将写入RESULTS_FOLDER/nnUNet文件夹。对于我们的项目来说，就是会存在/home/work/nnUNet/nnUNetFrame/DATASET/nnUNet_trained_models/nnUNet这个路径下。

每次训练都会获得一个自动生成的输出文件夹名称，根据我们的训练配置，我们会得到3d_fullres/Task001_BrainTumour这个文件夹。关于它的树状图如下（为简洁起见，有些文件仅在一个文件夹下详细展开）：

如果你的训练成功了，应该会得到和我下图一样的结果

下面详细讲讲这些训练后得到的文件都是什么

现在我们想看看我们训练的结果怎么样，有两种方法：

第一种是打开progress.png来从图像上直观的感受一下，它大概长这样：

第二种是打开validation_raw/summary.json，从里面我们不仅可以看到对每一个验证数据的评价，更可以在最末尾看到它们的平均值。

这里我们主要关注这个Dice分数，之前说到了脑瘤数据集分为四个模态，而labels的数量也是4（参见dataset.json）：

包含背景在内的4个标签，分别是背景、坏疽(NET,non-enhancing tumor)、浮肿区域(ED,peritumoral edema)、增强肿瘤区域(ET,enhancing tumor)，如下图，它们的平均dice分数约为0.9994，0.8770，0.7780，0.8728

八、确定最佳U-Net配置

本文只说明了3d_fullres的训练，完整的nnUNet流程还需要跑2d和3d级联的，然后进行三种的择优。不过从实际性能来说，一般3d级联≥3d＞2d，是否跑其他两种需要自己考虑。

训练完所有模型后，使用以下命令自动确定用于测试集预测的U-Net配置：

所有指定配置都需要完成所有的5折训练！

对于未配置级联的数据集，请改用“-m 2d 3d_fullres”。如果您只想探索配置的某些子集，可以使用“-m”命令指定。

所以说，如果你训练完所有模型，才可以使用这一步，如果没有，那可以直接推理

九、运行推理

1.准备测试集

让我们回到你刚刚做好数据集转换的那个数据集：/home/work/nnUNet/nnUNetFrame/DATASET/nnUNet_raw/nnUNet_raw_data/Task001_BrainTumour，在里面创建inferTs这个文件夹，用于存放待推理测试集的推理结果。然后，我会选择将原本的imagesTs重命名为imagesTs0，它是我们下载数据集时给我们的几十个测试集，然后新建一个imagesTs，里面只放一个测试集。（解释：这个道理其实很明显，推理太多数据集太久了，先用一个试试）