nnUNet框架小结（CUDA9.0+torch1.1\使用nrrd数据集）
nnUNet框架小结（CUDA9.0+torch1.1\使用nrrd数据集）
nnUNet框架效果很好，依靠一些技巧，将分割任务进行了大统一，并在很多任务上得到了非常好的成绩。该框架认为更多的提N升其实在于理解数据，并针对医学数据采用适当的预处理和训练方法。nnUNet可谓医学图像分割的大杀器。
本文将介绍在使用nnUNet途中的心路历程，帮助同需要使用该框架的小伙伴解决困惑。笔者使用的服务器的基础环境配置是CUDA9.1+torch1.1。但是现在github上的框架版本已经很新了，为了和笔者的环境搭配，便使用较为旧版本的nnUNet来进行使用。
一、 安装框架
-
首先服务器上已经存在pytorch1.1 + cuda9.0的基础环境。
-
安装NVIDIA-Apex：
这是英伟达的一个用于混合精度训练的插件，请不要直接pip，跟着下面的操作来：
第一步：打开Apex所在项目网站，往下拉便可以看到QuickStart，已经很详细。
第二步：在你用来安装环境的目录下打开终端，git clone https://github.com/NVIDIA/apex；
第三步：cd apex进入你刚才下载下来的apex文件夹里面
第四步：pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./【这步出现问题尝试使用python setup.py install --cuda_ext --cpp_ext，更多问题参考这里】 -
安装hiddenlayer
pip install --upgrade git+https://github.com/nanohanno/hiddenlayer.git@bugfix/get_trace_graph#egg=hiddenlayer(没有换行，这是一行代码) -
安装nnUNet
第一步：在home下创建nnUNet_sd_loss文件夹(自己定)，在这个文件夹内打开终端git clone https://github.com/MIC-DKFZ/nnUNet.git 第二步：
cd nnUNet 第三步：在nnUNet文件夹下 执行
git checkout 6ef1abe77625c0a72d4cfb8fd0b3b417ac00ef57 这一步为的是将nnUNet的版本回退到符合cuda9.0toch1.1的版本。如果该版本和你的环境不符合，你仍然通过在github上查历史版本从而进一步修改。
第四步：
pip install -e .（别忘了加 . ）
当你安装完成这些以后，你的每一次对nnUNet的操作，都会在命令行里以nnUNet_开头，代表着你的nnUNet开始工作的指令。
二、准备数据
1. 更改文件格式
这里使用的数据是nrrd格式的，所以我们需要先将nrrd格式的图像转换成符合框架要求的额nii.gz格式。
import os
from glob import glob
import nrrd #pip install pynrrd, if pynrrd is not already installed
import nibabel as nib #pip install nibabeModuleNotFoundError: No module named 'nrrd'l, if nibabel is not already installed
import numpy as np
baseDir = os.path.normpath('data')
outputDir = os.path.normpath('imagesTr')
files = glob(baseDir+'/*.nrrd')
for file in files:
#load nrrd
_nrrd = nrrd.read(file)
data = _nrrd[0]
header = _nrrd[1]
#save nifti
img = nib.Nifti1Image(data, np.eye(4))
filename=os.path.basename(file)
nib.save(img,os.path.join(outputDir, filename[-8:-5] + '.nii.gz'))
这里使用上述的脚本完成，完成之后我们就有nii.gz格式的原始数据集。
2. 创建文件夹目录结构
接下来在刚刚的nnUNet_sd_loss文件夹下建立DATASET文件夹。继而进入DATASET，继续创建nnUNet_preprocessed、nnUNet_raw、nnUNet_trained_models三个文件夹。第二个用来存放原始的你要训练的数据，第一个用来存放原始数据预处理之后的数据，第三个用来存放训练的结果。再继续进入上面第二个文件夹nnUNet_raw，创建nnUNet_cropped_data、nnUNet_raw_data两个文件夹，第二个为原始数据，第一个为crop以后的数据。
-
创建自己的任务文件夹
进入文件夹nnUNet_raw_data，创建一个名为Task08_ASOCA的文件夹（解释：这个Task08_ASOCA是nnUNet的子任务名，你可以对这个任务的数字ID进行任意的命名，比如你要分割心脏，你可以起名为Task01_Heart，比如你要分割肾脏，你可以起名为Task02_Kidney，前提是必须按照这种格式）
第五步：将刚刚通过脚本转换的数据集放在上面创建好的任务文件夹下，下面还以Task08_ASOCA为例，解释下数据应该怎么存放和编辑：-
你会发现目录是这个样子的：json文件是对三个文件夹内容的字典呈现（关乎你的训练），imagesTr是你的训练数据集，打开后你会发现很多的有序的nii.gz的训练文件，而labelsTr里时对应这个imagesTr的标签文件，同样为nii.gz。目前只能是nii.gz文件，nii文件都不行。训练阶段的imagesTs文件夹先不管，其实这个文件夹出现在任何位置都可以。（解释：nnUNet使用的是五折交叉验证，并没有验证集）
-
其中dataset.json，本来的nrrd数据集是没有，需要我们自己生成，在上面的目录结构下（nii.gz数据集文件已经正确放入各个文件夹中），执行python文件getJson.py，具体的代码如下，各位可以根据自己的数据集特定进行修改。
import glob
import os
import re
import json
from collections import OrderedDict
#将YOUR DIR替换成你自己的目录
path_originalData = ""
def list_sort_nicely(l):
""" Sort the given list in the way that humans expect.
"""
def tryint(s):
try:
return int(s)
except:
return s
def alphanum_key(s):
""" Turn a string into a list of string and number chunks.
"z23a" -> ["z", 23, "a"]
"""
return [ tryint(c) for c in re.split('([0-9]+)', s) ]
l.sort(key=alphanum_key)
return l
train_image = list_sort_nicely(glob.glob(path_originalData+"imagesTr/*"))
train_label = list_sort_nicely(glob.glob(path_originalData+"labelsTr/*"))
test_image = list_sort_nicely(glob.glob(path_originalData+"imagesTs/*"))
test_label = list_sort_nicely(glob.glob(path_originalData+"labelsTs/*"))
train_image = ["{}".format(patient_no.split('/')[-1]) for patient_no in train_image]
train_label = ["{}".format(patient_no.split('/')[-1]) for patient_no in train_label]
test_image = ["{}".format(patient_no.split('/')[-1]) for patient_no in test_image]
#输出一下目录的情况，看是否成功
print(len(train_image),len(train_label),len(test_image),len(test_label), train_image[0])
#####下面是创建json文件的内容
#可以根据你的数据集，修改里面的描述
json_dict = OrderedDict()
json_dict['name'] = "vessel"
json_dict['description'] = " Segmentation"
json_dict['tensorImageSize'] = "3D"
json_dict['reference'] = "see challenge website"
json_dict['licence'] = "see challenge website"
json_dict['release'] = "0.0"
#这里填入模态信息，0表示只有一个模态，还可以加入“1”：“MRI”之类的描述，详情请参考官方源码给出的示例
json_dict['modality'] = {
"0": "CT"
}
#这里为label文件中的多个标签，比如这里有血管、胆管、结石、肿块四个标签，名字可以按需要命名
json_dict['labels'] = {
"0": "Background",
"1": "vessel "#血管
}
#下面部分不需要修改>>>>>>
json_dict['numTraining'] = len(train_image)
json_dict['numTest'] = len(test_image)
json_dict['training'] = []
for idx in range(len(train_image)):
json_dict['training'].append({'image': "./imagesTr/%s" % train_image[idx], "label": "./labelsTr/%s" % train_label[idx]})
json_dict['test'] = ["./imagesTs/%s" % i for i in test_image]
with open(os.path.join(path_originalData, "dataset.json"), 'w') as f:
json.dump(json_dict, f, indent=4, sort_keys=True)
#<<<<<<<
一般情况下只需要注意模态信息，是ct还是核磁共振，以及labels有几类分别表示什么意思，此外注意好目录设置就行了。生成的json文件如下：
{
"description": " Segmentation",
"labels": {
"0": "Background",
"1": "vessel "
},
"licence": "see challenge website",
"modality": {
"0": "CT"
},
"name": "vessel",
"numTest": 20,
"numTraining": 40,
"reference": "see challenge website",
"release": "0.0",
"tensorImageSize": "3D",
"test": [
"./imagesTs/0.nii.gz",
"./imagesTs/1.nii.gz",
"./imagesTs/2.nii.gz",
"./imagesTs/3.nii.gz",
"./imagesTs/4.nii.gz",
"./imagesTs/5.nii.gz",
"./imagesTs/6.nii.gz",
"./imagesTs/7.nii.gz",
"./imagesTs/8.nii.gz",
"./imagesTs/9.nii.gz",
"./imagesTs/10.nii.gz",
"./imagesTs/11.nii.gz",
"./imagesTs/12.nii.gz",
"./imagesTs/13.nii.gz",
"./imagesTs/14.nii.gz",
"./imagesTs/15.nii.gz",
"./imagesTs/16.nii.gz",
"./imagesTs/17.nii.gz",
"./imagesTs/18.nii.gz",
"./imagesTs/19.nii.gz"
],
"training": [
{
"image": "./imagesTr/0.nii.gz",
"label": "./labelsTr/0.nii.gz"
},
{
"image": "./imagesTr/1.nii.gz",
"label": "./labelsTr/1.nii.gz"
},
{
"image": "./imagesTr/2.nii.gz",
"label": "./labelsTr/2.nii.gz"
},
{
"image": "./imagesTr/3.nii.gz",
"label": "./labelsTr/3.nii.gz"
},
{
"image": "./imagesTr/4.nii.gz",
"label": "./labelsTr/4.nii.gz"
},
{
"image": "./imagesTr/5.nii.gz",
"label": "./labelsTr/5.nii.gz"
},
{
"image": "./imagesTr/6.nii.gz",
"label": "./labelsTr/6.nii.gz"
},
{
"image": "./imagesTr/7.nii.gz",
"label": "./labelsTr/7.nii.gz"
},
{
"image": "./imagesTr/8.nii.gz",
"label": "./labelsTr/8.nii.gz"
},
{
"image": "./imagesTr/9.nii.gz",
"label": "./labelsTr/9.nii.gz"
},
{
"image": "./imagesTr/10.nii.gz",
"label": "./labelsTr/10.nii.gz"
},
{
"image": "./imagesTr/11.nii.gz",
"label": "./labelsTr/11.nii.gz"
},
{
"image": "./imagesTr/12.nii.gz",
"label": "./labelsTr/12.nii.gz"
},
{
"image": "./imagesTr/13.nii.gz",
"label": "./labelsTr/13.nii.gz"
},
{
"image": "./imagesTr/14.nii.gz",
"label": "./labelsTr/14.nii.gz"
},
{
"image": "./imagesTr/15.nii.gz",
"label": "./labelsTr/15.nii.gz"
},
{
"image": "./imagesTr/16.nii.gz",
"label": "./labelsTr/16.nii.gz"
},
{
"image": "./imagesTr/17.nii.gz",
"label": "./labelsTr/17.nii.gz"
},
{
"image": "./imagesTr/18.nii.gz",
"label": "./labelsTr/18.nii.gz"
},
{
"image": "./imagesTr/19.nii.gz",
"label": "./labelsTr/19.nii.gz"
},
{
"image": "./imagesTr/20.nii.gz",
"label": "./labelsTr/20.nii.gz"
},
{
"image": "./imagesTr/21.nii.gz",
"label": "./labelsTr/21.nii.gz"
},
{
"image": "./imagesTr/22.nii.gz",
"label": "./labelsTr/22.nii.gz"
},
{
"image": "./imagesTr/23.nii.gz",
"label": "./labelsTr/23.nii.gz"
},
{
"image": "./imagesTr/24.nii.gz",
"label": "./labelsTr/24.nii.gz"
},
{
"image": "./imagesTr/25.nii.gz",
"label": "./labelsTr/25.nii.gz"
},
{
"image": "./imagesTr/26.nii.gz",
"label": "./labelsTr/26.nii.gz"
},
{
"image": "./imagesTr/27.nii.gz",
"label": "./labelsTr/27.nii.gz"
},
{
"image": "./imagesTr/28.nii.gz",
"label": "./labelsTr/28.nii.gz"
},
{
"image": "./imagesTr/29.nii.gz",
"label": "./labelsTr/29.nii.gz"
},
{
"image": "./imagesTr/30.nii.gz",
"label": "./labelsTr/30.nii.gz"
},
{
"image": "./imagesTr/31.nii.gz",
"label": "./labelsTr/31.nii.gz"
},
{
"image": "./imagesTr/32.nii.gz",
"label": "./labelsTr/32.nii.gz"
},
{
"image": "./imagesTr/33.nii.gz",
"label": "./labelsTr/33.nii.gz"
},
{
"image": "./imagesTr/34.nii.gz",
"label": "./labelsTr/34.nii.gz"
},
{
"image": "./imagesTr/35.nii.gz",
"label": "./labelsTr/35.nii.gz"
},
{
"image": "./imagesTr/36.nii.gz",
"label": "./labelsTr/36.nii.gz"
},
{
"image": "./imagesTr/37.nii.gz",
"label": "./labelsTr/37.nii.gz"
},
{
"image": "./imagesTr/38.nii.gz",
"label": "./labelsTr/38.nii.gz"
},
{
"image": "./imagesTr/39.nii.gz",
"label": "./labelsTr/39.nii.gz"
}
]
}
-
设置nnUNet读取文件的路径
要让nnUNet知道你的文件存放在哪儿需要在环境中创建一个路径。
找到.bashrc文件，打开；在文档末尾添加下面三行，右上角保存文件，观察下面保存成功后关闭。export nnUNet_raw_data_base="/home/hongqq/nnUNet_sd_loss/DATASET/nnUNet_raw" export nnUNet_preprocessed="/home/hongqq/nnUNet_sd_loss/DATASET/nnUNet_preprocessed" export RESULTS_FOLDER="/home/hongqq/nnUNet_sd_loss/DATASET/nnUNet_trained_models"在home下打开终端，输入
source .bashrc来更新该文档。nnUNet已经知道怎么读取你的文件了。
三、训练模型
-
转化数据集，使得其能够被框架识别
nnUNet_convert_decathlon_task -i /home/hongqq/nnUNet_sd_loss/DATASET/nnUNet_raw/nnUNet_raw_data/Task08_ASOCA转换之后会发现，在这个Task08_ASOCA文件夹旁边多了一个Task008_ASOCA，里面的文件除了labels，其他末尾都多了_0000，这就是你的数据格式是否正确的标志。。
-
数据预处理
nnUNet_plan_and_preprocess -t 8上面的参数8是你的任务id。
-
开始训练
-
单卡训练：执行
nnUNet_train 3d_fullres nnUNetTrainerV2 8 4
注意：默认在第一块gpu（索引为0）上进行训练，如果想指定某个gpu，请先执行：
export CUDA_VISIBLE_DEVICES=X，X为你指定的gpu索引。再执行上面的命令。8代表你的任务ID，4代表五折交叉验证中的第4折（0代表分成五折后的第一折）。具体的参数，可以去翻一翻源代码就知道了。
-
多卡训练：
比如我现在要在0和1两张卡上执行训练：
-
先执行
export CUDA_VISIBLE_DEVICES=0,1 -
再执行
nnUNet_train_DP 3d_fullres nnUNetTrainerV2_DP 8 4 -gpus 2多卡并不显著提高训练速度，但是可以应对某些情况下单卡显存不足的情况。
-
我们这里对上图的输出信息进行一些说明：
- 1 loss的最优值是趋近于-1的，因为默认的loss函数是交叉熵与softdice的混合，其最优值为-1。
- 2 此处是验证的dice值。
- 3 梯度溢出后，框架会自动跳过。
- 4 每个epoch 正常会需要10多分钟的时间。
4.训练结束时也许会报错：
如图当所有的epoch都跑完了的时候，报source tensor must be contiguous.的错误，解决方法是修改下to_torch.py源代码：
训练结束可以在文件夹中找到loss的下降图：
四、利用训练结果进行预测
- 创建文件夹
创建两个空文件夹inferTs、labelsTs使你的Task008文件底下像这样：
labelsTs中存放了测试集的标签，inferTs是我待推理测试集的推理结果。
-
进行预测
nnUNet_predict -i /home/你的主机用户名/nnUNet_sd_loss/DATASET/nnUNet_raw/nnUNet_raw_data/Task008_ASOCA/imagesTs/ -o /home/你的主机用户名/nnUNet_sd_loss/DATASET/nnUNet_raw/nnUNet_raw_data/Task008_ASOCA/inferTs -t 8 -m 3d_fullres -f 4nnUNet_predict：执行预测的命令；
-i: 输入（你的待推理测试集）；
-o: 输出（测试集的推理结果）；
-t: 你的任务对应的数字ID；
-m: 对应的训练时使用的网络架构；
-f: 数字4代表使用五折交叉验证训练出的模型；注意事项：如果该训练类已经预测过一次，再次训练进行预测时，记得将infer文件夹中的数据删除，否则无法成功预测。
-
将预测结果转回nrrd
这里我们使用一个软件——Slicer，搜索3D Slicer下载安装稳定版本即可，然后就可以使用该软件将nii.gz转化成nrrd。
仅利用导入data，和保存data就可以进行数据格式的转化，可以进行批量操作。
五、 loss函数的修改
loss函数的修改只需要对某个基类进行继承修改初始函数即可，当然修改的地方要对，不然的话运行之后框架无法识别你写的新类，就会报错。如下：
正确修改的目录为/home/hongqq/nnUNet_sd_loss/nnUNet/nnunet/training/network_training/，在此目录下新建一个类，然后继承nnUNetTrainerV2，改写其初始化函数，指定新的loss函数。
from nnunet.training.network_training.nnUNetTrainerV2 import nnUNetTrainerV2
from nnunet.training.loss_functions.dice_loss import Dice_and_Cl_loss
from nnunet.utilities.nd_softmax import softmax_helper
class nnUNetTrainerV2_DiceClDice(nnUNetTrainerV2):
def __init__(self, plans_file, fold, output_folder=None, dataset_directory=None, batch_dice=True, stage=None,
unpack_data=True, deterministic=True, fp16=False):
super().__init__(plans_file, fold, output_folder, dataset_directory, batch_dice, stage, unpack_data,
deterministic, fp16)
self.loss = Dice_and_Cl_loss({'batch_dice': self.batch_dice, 'smooth': 1e-5, 'do_bg': False})
self.max_num_epochs=500
当然如果需要进行双卡训练，则需要对nnUNetTrainerV2_DP进行继承改写初始化函数。
from nnunet.training.network_training.nnUNetTrainerV2_DP import nnUNetTrainerV2_DP
from nnunet.training.loss_functions.dice_loss import Dice_and_softcl_loss
class nnUNetTrainerV2_DP_DiceSoftClDice(nnUNetTrainerV2_DP):
def __init__(self, plans_file, fold, output_folder=None, dataset_directory=None, batch_dice=True, stage=None,
unpack_data=True, deterministic=True, num_gpus=1, distribute_batch_size=False, fp16=False):
super(nnUNetTrainerV2_DP, self).__init__(plans_file, fold, output_folder, dataset_directory, batch_dice, stage,
unpack_data, deterministic, fp16)
self.init_args = (plans_file, fold, output_folder, dataset_directory, batch_dice, stage, unpack_data,
deterministic, num_gpus, distribute_batch_size, fp16)
self.num_gpus = num_gpus
self.distribute_batch_size = distribute_batch_size
self.dice_smooth = 1e-5
self.dice_do_BG = False
self.loss_weights = None
self.loss = Dice_and_softcl_loss({'batch_dice': self.batch_dice, 'smooth': 1e-5, 'do_bg': False})
self.max_num_epochs=500
然后就可以使用改好的类进行训练了，以上面改的两个类为例子：
-
单卡训练：执行
nnUNet_train 3d_fullres nnUNetTrainerV2_DiceClDice 8 4 -
多卡训练：
在0和1两张卡上执行训练：
- 先执行
export CUDA_VISIBLE_DEVICES=0,1 - 再执行
nnUNet_train_DP 3d_fullres nnUNetTrainerV2_DP_DiceSoftClDice 8 4 -gpus 2
- 先执行
此时使用训练结果进行预测时，需要在预测语句中加上-tr 你自定义的Trainer，不然的话框架会找不到你保存的模型在哪里。
六，预测结果可视化
大多数医学图像的分割都是有具体意义的，比如对血管的分割。所有我们有时需要对自己的预测结果进行可视化，以便于对结果有个大致的分析。
import matplotlib.pyplot as plt
import numpy as np
import nrrd
import os
from glob import glob
def plotNrrd(filepath):
arr0=nrrd.read('C:/Users/32920/Desktop/nnUnet结果/'+filepath)[0]
x=[]
y=[]
z=[]
for i in range(len(arr0)):
for j in range(len(arr0[0])):
for k in range(len(arr0[0][0])):
if(arr0[i][j][k]==1):
print(i+1,' ',j+1,' ',k+1)
x.append(i+1)
y.append(j+1)
z.append(k+1)
fileindex=(filepath.split('/')[-1]).split('.')[0]
with open('nrrd'+fileindex+'_x.txt','a')as f:
f.write(','.join(map(str,x)))
with open('nrrd'+fileindex+'_y.txt','a')as f1:
f1.write(','.join(map(str,y)))
with open('nrrd'+fileindex+'_z.txt','a')as f2:
f2.write(','.join(map(str,z)))
ax = plt.figure().add_subplot(111, projection = '3d')
#基于ax变量绘制三维图
#xs表示x方向的变量
#ys表示y方向的变量
#zs表示z方向的变量，这三个方向上的变量都可以用list的形式表示
#m表示点的形式，o是圆形的点，^是三角形（marker)
#c表示颜色（color for short）
ax.scatter(x, y, z, c = 'r', marker = '^') #点为红色三角形
#设置坐标轴
ax.set_xlabel('X Label')
ax.set_ylabel('Y Label')
ax.set_zlabel('Z Label')
plt.title(fileindex+'.nrrd')
#显示图像
plt.savefig(filepath+'.png')
if __name__=='__main__':
# plotNpz('Predict_Masks_labeled_org/30.npz')
#plotNrrd('ASOCA/Test_Masks/32.nrrd')
files = glob('*.nrrd')
for file in files:
plotNrrd(file)
得到可视化结果如图所示：
参考博主：https://blog.csdn.net/weixin_42061636/article/details/107623757