Ubuntu下使用nnUNet 训练自己的数据集（只管实现，不讲原理，通俗易懂）

方便日后跑通nnUNet进行对比实验，写下本文。

本文只管如何在自己的数据集上跑通nnUNet，不讲解原理，因为我也是参考别人博文来做的，详情可参看Tina的四篇博文

一、什么是nnUNet？

nnU-Net是由德国癌症研究中心、海德堡大学以及海德堡大学医院研究人员(Fabian Isensee, Jens Petersen, Andre Klein)提出来的一个自适应任何新数据集的医学影像分割框架，该框架能根据给定数据集的属性自动调整所有超参数，整个过程无需人工干预。通过对数据的自适应处理，能在多个挑战赛中取得第一名，足以见得这个框架是有多么优秀。时至今日，nnUNet仍能凭借最为简单的U-Net结构在六个得到公认的分割挑战中实现了最先进的性能，在多个医学图像分割比赛中长期处于霸榜地位。

二、下载并配置nnUNet环境

1、下载，直接通过下面命令下载即可

pip install nnunet

下载好之后能看到根目录下多了一个文件夹，如下：

该目录下长这样的：

2、配置环境变量

在主目录中的.bashrc文件中设置路径，有部分朋友找不到该文件的可以在文件管理器里面按Ctrl + H 即可显示.bashrc文件。

随后直接对.bashrc文件进行编辑，在文件底部加入一下内容（有的在安装nnUNet时就自动生成了）

export nnUNet_raw_data_base="/home/gigabyte/nnUNetYV/DATASET/nnUNet_raw"
export nnUNet_preprocessed="/home/gigabyte/nnUNetYV/DATASET/nnUNet_preprocessed"
export RESULTS_FOLDER="/home/gigabyte/nnUNetYV/DATASET/nnUNet_trained_models"

之后nnUNet在运行是会自动找到这三个路径进行相关的读写操作。

三、将自己的数据集处理成nnUNet能识别的格式

1、文件夹结构安排

nnUNet需要结构化格式的数据集，并将原始数据存放在nnUNet_raw_data_base/nnUNet_raw_data文件夹下。其中，文件夹一定要命名为（“Task”“ID”_“任务名称”)， “ID”是三位数的整, 任务名称可以是数据集的简写。

任务的文件夹名称如下，下面分别是5个分割任务：（引用自Tina）

再比如本人的分割任务： Task676_NPC，以’NPC’作为任务名称任务ID为676。

其中，每个文件（夹）的含义如下：

• imagesTr: 训练图像。
• imagesTs（可选，上图中我的就没有测试集）: 测试图像
• labelTr: 训练集标签。
• dataset.json : 包含数据集的元数据， 如任务名字，模态，标签含义，训练集包含的图像地址

2、图像文件的命名要求

使用nnUNet时，对图像的名称是有要求的！图像和标签的名字具有一定的关联性，所有图像（包括标签）必须是3D nifti文件（.nii.gz）！

图像的统一命名格式case_identifier_XXXX.nii.gz， 标签的统一命名格式case_identifier.nii.gz

XXXX是模态标识符, 四位整数。

此命名方案将导致以下文件夹结构，必须严格将其数据转换为这种格式，参考Tina的博文如下：

 最后，附上本人任务（Task676_NPC）的文件夹示例：

 3、Json文件内容解读

一个合格的Json文件应该包含如下信息：

name: 数据集名字
dexcription: 对数据集的描述
modality: 模态，可以是CT或者MRI数据，nnU-Net会根据不同模态进行不同的预处理
labels: label中，不同的数值代表的类别
numTraining: 训练集数量
numTest: 测试集数量
training: 训练集的image 和 label 地址对
test： 只包含测试集的image. 这里跟Training不一样

引用自Tina的图片：

 本人任务（Task676_NPC）的数据集对应的Json文件如下：

 现在我们知道长什么样了，那么怎么把我们的数据准备成这个样子呢。最简单的方法就是copy别人的代码，在前人的基础上修改一下。如： dataset_conversion

—————————————————2022.09.04 更新开始————————————————

再次需要用nnUNet跑LiTS数据集，附上本人生成json文件的代码如下：

#    Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
from collections import OrderedDict

import json
import os

if __name__ == "__main__":

    train_dir = r"D:\Doctor_Li_Brest_Data\LITS17\LiTS_nnUNet_dataset\labelsTr"
    test_dir =r'D:\Doctor_Li_Brest_Data\LITS17\LiTS_nnUNet_dataset\imagesTs'
    output_folder = r"D:\Doctor_Li_Brest_Data\LITS17\LiTS_nnUNet_dataset"

    json_dict = OrderedDict()
    json_dict['name'] = "LiTS"
    json_dict['description'] = "LiTS"
    json_dict['tensorImageSize'] = "3D"
    json_dict['reference'] = "see challenge website"
    json_dict['licence'] = "see challenge website"
    json_dict['release'] = "0.0"
    json_dict['modality'] = {
        "0": "CT"
        # "0": "MR"
    }

    json_dict['labels'] = {
        "0": "background",
        # "1": "liver",
        "1": "tumor"
    }

    file_name_list = os.listdir(train_dir)
    file_id_list = []
    for name in file_name_list:
        file_id_list.append(name.split('.nii.gz')[0])

    file_name_list = os.listdir(test_dir)
    file_name_test_ids = []
    for name in file_name_list:
        file_name_test_ids.append(name.split('_')[0])


    json_dict['numTraining'] = len(file_id_list)
    json_dict['training'] = [{'image': "./imagesTr/%s.nii.gz" % i, "label": "./labelsTr/%s.nii.gz" % i} for i in file_id_list]
    json_dict['numTest'] = 0
    json_dict['test'] = []
    # json_dict['numTest'] = len(file_name_test_ids)  ##有测试集就解开注释
    # json_dict['test'] = ["./imagesTs/%s.nii.gz" % i for i in file_name_test_ids]  ##有测试集就解开注释

    with open(os.path.join(output_folder, "dataset.json"), 'w') as f:
        json.dump(json_dict, f, indent=4, sort_keys=True)

—————————————————2022.09.04 更新结束————————————————

至此，数据集配置完成。

四、nnUNet模型跑起来

1、nnUNet框架对数据进行自适应地预处理

首先，当然是激活nnUNet的环境啦！

conda activate nnUNet***environmental

使用命令： nnUNet_plan_and_preprocess

nnUNet_plan_and_preprocess -t XXX --verify_dataset_integrity

作用： 它会自动读取 数据集的属性，例如图像大小、体素间距、强度信息等。并转换成模型需要的格式。

XXX是与您的任务名称相关联的整数标识符TaskXXX_XXX（如本人的任务Task676_NPC， 就可以写676）

即XXX是任务号码，如676.

–verify_dataset_integrity 对数据集执行一些检查以确保它与 nnU-Net 兼容。如果此检查已通过一次，则可以在以后的运行中省略。

 Task676_NPC任务预处理完成后的结果保存如下文件夹：

—————————————————2022.09.04 更新开始————————————————

LiTS的处理过程截图如下：

LiTS预处理后的结果保存如下：

 —————————————————2022.09.04 更新结束————————————————

2、模型训练

使用命令 nnUNet_train

该命令参数很多，可以使用nnUNet_train -h查看参数的含义

该命令的一般结构为：

nnUNet_train CONFIGURATION TRAINER_CLASS_NAME TASK_NAME_OR_ID FOLD (additional options)

CONFIGURATION： 模型架构，三种Unet: 2D U-Net, 3D U-Net and a U-Net Cascade(U-Net级联)。
TRAINER_CLASS_NAME： 使用的model trainer. 默认为nnUNetTrainerV2就可以
TASK_NAME_OR_ID： 任务全名TaskXXX_MYTASK或者是ID号，如Task676_NPC或676
FOLD： 第几折交叉验证，可选 [0, 1, 2, 3, 4]，一共五折。

举例一： 2D U-Net 训练第三折

nnUNet_train 2d nnUNetTrainerV2 676 3

举例二： 3D full resolution U-Net 训练第二折

nnUNet_train 3d_fullres nnUNetTrainerV2 676 2

训练过程中，终端会输出以下信息：大致包括使用的模型，数据信息，损失函数，网络架构，每个epoch的输出。

每次训练都会获得一个自动生成的输出文件夹名称：

 其中，你想要查看的信息都在日志文件training_log文件中；

debug.json：包含用于训练此模型的蓝图和推断参数的摘要。不容易阅读，但对调试非常有用;
model_best.model / model_best.model.pkl：训练期间确定的最佳模型的检查点文件。
model_final_checkpoint.model / model_final_checkpoint.model.pkl：最终模型的检查点文件，用于验证和推理的内容。
progress.png：训练过程中的训练（蓝色）和验证（红色）损失图。还显示了评估指标的近似值（绿色）。这个近似值是前景类的平均 Dice 分数。

—————————————————2022.09.04 更新开始————————————————

五、修改训练的Epoch

由于nnUNet默认的epoch=1000，训练是在太久了，可用以下方式修改训练的epoch的默认值

1、在nnUNet的源码目录中找到nnUNetTrainerV2.py，该文件的位置如下：

2、直接将原来的self.max_num_epochs=1000 改为 自己想要的值，如改为23

部分细节没写到位，有疑问可以私信，尽量回复。

码字不易，有参考意义就给个赞呗，谢谢~~