nnU-Net，它是语义图像分割中一个非常强大的基线。

在本文中，您将了解如下：

1. 对nnU-Net的主要贡献进行简明概述。

2. 了解如何将nnU-Net应用于您自己的数据集。

nnU-Net简史

作为公认的最先进的图像分割模型，nnU-Net在2D和3D图像处理方面都是一股不屈不挠的力量。它的性能非常强大，可以作为新的计算机视觉体系结构的基准。

这个强大的工具基于U-Net模型， U-Net于2015年首次亮相。名称“nnU-Net”代表“No New U-Net”，这是对其设计没有引入革命性建筑变化这一事实的认可。相反，它利用现有的U-Net结构，通过一套巧妙的优化策略，充分发挥其潜力。

与许多现代神经网络不同，nnU-Net不依赖于残差连接、密集连接或注意机制。它的优势在于其细致的优化策略，其中包括重采样、归一化、损失函数的明智选择、优化器设置、数据增强、基于补丁的推理和跨模型集成等技术。这种整体的方法允许nnU-Net突破原有的U-Net架构所能达到的极限。

探索nnU-Net中的多样化架构

虽然它看起来像是一个单一的实体，但nnU-Net实际上是三种不同类型的u - net的总称：

1. 2D U-Net：可以说是最著名的变种，它直接操作2D图像。

2. 3D U-Net：这是2D U-Net的扩展，能够通过应用3D卷积直接处理3D图像。

3. U-Net Cascade：该模型生成低分辨率的分割，并随后对其进行细化。

每种体系结构都有其独特的优势，并且不可避免地有一定的局限性。

例如，使用2D U-Net进行3D图像分割可能看起来违反直觉，但在实践中，它仍然可以非常有效。这是通过将3D体切割成2D平面来实现的。

虽然3D U-Net可能看起来更复杂，但考虑到其更高的参数数量，它并不总是最有效的解决方案。特别是，3D U-Net经常与各向异性作斗争，当空间分辨率沿着不同的轴(例如，沿着x轴1mm和沿着z轴1.2 mm)不同时，就会出现各向异性。

当处理大尺寸图像时，U-Net级联变体变得特别方便。它采用一个初步的模型来压缩图像，然后是一个标准的3D U-Net，输出低分辨率的分割。然后，生成的预测被放大，从而产生一个精细的、全面的输出。

通常，该方法涉及在nnU-Net框架内训练所有三种模型变体。接下来的步骤可能是在三人中选择表现最好的或采用集成技术。其中一项技术可能涉及整合2D和3D U-Nets的预测。

然而，值得注意的是，这个过程可能非常耗时(也很花钱，因为你需要GPU积分)。如果您的约束只允许训练单个模型，请不要担心。您可以选择只训练一个模型，因为集成模型只带来非常小的收益。

下表说明了与特定数据集相关的最佳表现模型变体：

网络拓扑的动态适应

考虑到图像大小的显著差异(考虑肝脏图像的中位数形状为482 × 512 × 512，而海马图像的中位数形状为36 × 50 × 35)， nnU-Net智能地适应输入补丁大小和每轴池化操作的数量。这本质上意味着自动调整每个数据集的卷积层数，促进空间信息的有效聚合。除了适应不同的图像几何形状外，该模型还考虑了技术限制，例如可用内存。

值得注意的是，该模型并不直接对整个图像进行分割，而是对精心提取的具有重叠区域的补丁进行分割。随后对这些补丁的预测进行平均，从而得到最终的分割输出。

但是拥有大的补丁意味着更多的内存使用，并且批处理大小也消耗内存。所采取的权衡是始终优先考虑补丁大小(模型的容量)而不是批处理大小(仅对优化有用)。

下面是用于计算最优补丁大小和批大小的启发式算法：

这是不同数据集和输入维度的情况：

现在让我们回顾一下nnU-Net中使用的所有技术：

训练

所有模型都是从头开始训练的，并在训练集上使用五倍交叉验证进行评估，这意味着原始训练数据集被随机分为五个相等的部分，或“折叠”。在这个交叉验证过程中，其中四个折叠用于模型的训练，剩下的一个折叠用于评估或测试。这个过程然后重复五次，五次中的每一次都被用作评估集。

对于损失，我们使用骰子和交叉熵损失的组合。这是图像分割中非常常见的损失。

数据增强技术

nnU-Net有非常强大的数据增强管道。作者使用了随机旋转、随机缩放、随机弹性变形、伽马校正和镜像。

基于补丁的推理

正如我们所说，该模型不是直接在全分辨率图像上进行预测，而是在提取的斑块上进行预测，然后将预测汇总。

它看起来是这样的：

注意：图片中间的补丁比旁边的补丁权重更大，因为它们包含更多的信息，模型对它们的处理效果更好

两两模型集成

如果你还记得的话，我们可以训练3种不同的模型，2D, 3D和级联。但是当我们进行推理时我们一次只能使用一个模型，对吧?

事实证明，不同的模型有不同的优点和缺点。因此，我们实际上可以将几个模型的预测结合起来，如果一个模型非常自信，我们就会优先考虑它的预测。

nnU-Net测试了3个可用模型中的2个模型的每种组合，并挑选出最好的一个。

实际上，有两种方法可以做到这一点：

1. 硬投票：对于每个像素，我们查看2个模型输出的所有概率，并选择概率最高的类。

2. 软投票：对于每个像素，我们对模型的概率取平均值，然后取最大概率的类。

实践

如果你不明白第一部分，不用担心，这是实践部分，你只需要跟着我，你仍然会得到最好的结果。

你需要GPU来训练模型，否则模型无法工作。你既可以在本地做，也可以在Google Collab上做，别忘了修改运行时>GPU

首先，你需要有一个数据集准备好输入图像和它们相应的分割。你可以按照我的教程下载这个现成的3D大脑分割数据集，然后你可以用你自己的数据集替换它。

下载数据

首先，您应该下载数据并将其放入data文件夹中，将包含分割的两个文件夹命名为“input”和“ground_truth”。

在本教程的其余部分，我将使用MindBoggle数据集进行图像分割。

我们得到了大脑的3D核磁共振扫描，我们想要分割白质和灰质：

它应该是这样的：

设置主目录

如果你在Google Colab上运行这个，设置collab = True，否则设置collab = False

collab = True
import os
import shutil
#libraries
from collections import OrderedDict
import json
import numpy as np
#visualization of the dataset
import matplotlib.pyplot as plt
import nibabel as nib
if collab:
    from google.colab import drive
    drive.flush_and_unmount()
    drive.mount('/content/drive', force_remount=True)
    # Change "neurosciences-segmentation" to the name of your project folder
    root_dir = "/content/drive/MyDrive/neurosciences-segmentation"

else:
    # get the dir of the parent dir
    root_dir = os.getcwd()
input_dir = os.path.join(root_dir, 'data/input')
segmentation_dir = os.path.join(root_dir, 'data/ground_truth')
my_nnunet_dir = os.path.join(root_dir,'my_nnunet')
print(my_nnunet_dir)

现在我们要定义一个为我们创建文件夹的函数：

def make_if_dont_exist(folder_path,overwrite=False):
    """
    creates a folder if it does not exists
    input:
    folder_path : relative path of the folder which needs to be created
    over_write :(default: False) if True overwrite the existing folder
    """
    if os.path.exists(folder_path):
        if not overwrite:
            print(f'{folder_path} exists.')
        else:
            print(f"{folder_path} overwritten")
            shutil.rmtree(folder_path)
            os.makedirs(folder_path)
    else:
      os.makedirs(folder_path)
      print(f"{folder_path} created!")

我们使用这个函数来创建“my_nnunet”文件夹，所有内容都将保存在这里：

os.chdir (root_dir)
make_if_dont_exist (my_nnunet,覆盖= False)
os.chdir(“my_nnunet”)
print(f"当前工作目录:{os.getcwd()}")

库安装

现在我们要安装所有的需求。首先，让我们安装nnunet库。如果你用的是笔记本，在一个单元格里运行：

!pip install nnunet

否则，您可以直接从终端安装nnunet

pip install nnunet

现在我们要克隆nnUnet git库和NVIDIA apex。这包含训练脚本以及GPU加速器。

!git clone https://github.com/MIC-DKFZ/nnUNet.git
!git clone https://github.com/NVIDIA/apex
# repository dir is the path of the github folder
respository_dir = os.path.join(my_nnunet_dir,'nnUNet')
os.chdir(respository_dir)
!pip install -e
!pip install --upgrade git+https://github.com/nanohanno/hiddenlayer.git@bugfix/get_trace_graph#egg=hiddenlayer

创建文件夹

nnUnet需要非常具体的文件夹结构。

task_name = 'Task001' #change here for different task name
# We define all the necessary paths
nnunet_dir = "nnUNet/nnunet/nnUNet_raw_data_base/nnUNet_raw_data"
task_folder_name = os.path.join(nnunet_dir,task_name) 
train_image_dir = os.path.join(task_folder_name,'imagesTr') # path to training images
train_label_dir = os.path.join(task_folder_name,'labelsTr') # path to training labels
test_dir = os.path.join(task_folder_name,'imagesTs') # path to test images
main_dir = os.path.join(my_nnunet_dir,'nnUNet/nnunet') # path to main directory
trained_model_dir = os.path.join(main_dir, 'nnUNet_trained_models') # path to trained models

最初，nnU-Net是为十项全能的不同任务挑战而设计的。如果您有不同的任务，只需为所有任务运行此单元格。

# Creation of all the folders
overwrite = False # Set this to True if you want to overwrite the folders
make_if_dont_exist(task_folder_name,overwrite = overwrite)
make_if_dont_exist(train_image_dir, overwrite = overwrite)
make_if_dont_exist(train_label_dir, overwrite = overwrite)
make_if_dont_exist(test_dir,overwrite= overwrite)
make_if_dont_exist(trained_model_dir, overwrite=overwrite)

你现在应该有一个这样的结构:

设置环境变量

脚本需要知道将raw_data放在哪里，可以在哪里找到预处理数据，以及必须在哪里保存结果。

os.environ['nnUNet_raw_data_base'] = os.path.join(main_dir,'nnUNet_raw_data_base')
os.environ['nnUNet_preprocessed'] = os.path.join(main_dir,'preprocessed')
os.environ['RESULTS_FOLDER'] = trained_model_dir

将文件移动到正确的存储库中:

我们定义了一个函数，将图像移动到nnunet文件夹中正确的存储库中:

def copy_and_rename(old_location,old_file_name,new_location,new_filename,delete_original = False):
    shutil.copy(os.path.join(old_location,old_file_name),new_location)
    os.rename(os.path.join(new_location,old_file_name),os.path.join(new_location,new_filename))
    if delete_original:
        os.remove(os.path.join(old_location,old_file_name))

现在让我们对输入图像和真实图像运行这个函数:

list_of_all_files = os.listdir(segmentation_dir)
list_of_all_files = [file_name for file_name in list_of_all_files if file_name.endswith('.nii.gz')]

for file_name in list_of_all_files:
    copy_and_rename(input_dir,file_name,train_image_dir,file_name)
    copy_and_rename(segmentation_dir,file_name,train_label_dir,file_name)

现在我们必须重命名要被nnUnet格式接受的文件，例如subject.nii.gz将变成subject_0000.nii.gz

def check_modality(filename):
    """
    check for the existence of modality
    return False if modality is not found else True
    """
    end = filename.find('.nii.gz')
    modality = filename[end-4:end]
    for mod in modality:
        if not(ord(mod)>=48 and ord(mod)<=57): #if not in 0 to 9 digits
            return False
    return True
def rename_for_single_modality(directory):
    for file in os.listdir(directory):
        if check_modality(file)==False:
            new_name = file[:file.find('.nii.gz')]+"_0000.nii.gz"
            os.rename(os.path.join(directory,file),os.path.join(directory,new_name))
            print(f"Renamed to {new_name}")
        else:
            print(f"Modality present: {file}")
rename_for_single_modality(train_image_dir)
# rename_for_single_modality(test_dir)

设置JSON文件

你主要需要修改两件事:

1. 模态(如果是CT或MRI，这会改变归一化)

2. 标签：输入您自己的类

overwrite_json_file = True #make it True if you want to overwrite the dataset.json file in Task_folder
json_file_exist = False
if os.path.exists(os.path.join(task_folder_name,'dataset.json')):
    print('dataset.json already exist!')
    json_file_exist = True
if json_file_exist==False or overwrite_json_file:
    json_dict = OrderedDict()
    json_dict['name'] = task_name
    json_dict['description'] = "Segmentation of T1 Scans from MindBoggle"
    json_dict['tensorImageSize'] = "3D"
    json_dict['reference'] = "see challenge website"
    json_dict['licence'] = "see challenge website"
    json_dict['release'] = "0.0"
    ######################## MODIFY THIS ########################
    #you may mention more than one modality
    json_dict['modality'] = {
        "0": "MRI"
    }
    #labels+1 should be mentioned for all the labels in the dataset
    json_dict['labels'] = {
        "0": "Non Brain",
        "1": "Cortical gray matter",
        "2": "Cortical White matter",
        "3" : "Cerebellum gray ",
        "4" : "Cerebellum white"
    }
    #############################################################

    train_ids = os.listdir(train_label_dir)
    test_ids = os.listdir(test_dir)
    json_dict['numTraining'] = len(train_ids)
    json_dict['numTest'] = len(test_ids)
    #no modality in train image and labels in dataset.json
    json_dict['training'] = [{'image': "./imagesTr/%s" % i, "label": "./labelsTr/%s" % i} for i in train_ids]
    #removing the modality from test image name to be saved in dataset.json
    json_dict['test'] = ["./imagesTs/%s" % (i[:i.find("_0000")]+'.nii.gz') for i in test_ids]
    with open(os.path.join(task_folder_name,"dataset.json"), 'w') as f:
        json.dump(json_dict, f, indent=4, sort_keys=True)
    if os.path.exists(os.path.join(task_folder_name,'dataset.json')):
        if json_file_exist==False:
            print('dataset.json created!')
        else:
            print('dataset.json overwritten!')

将数据预处理为nnU-Net格式

这将创建nnU-Net格式的数据集

# -t 1 means "Task001", if you have a different task change it
!nnUNet_plan_and_preprocess -t 1 --verify_dataset_integrity

训练模型

要训练3D U-Net:

#train 3D full resolution U net
!nnUNet_train 3d_fullres nnUNetTrainerV2 1 0 --npz 

要训练2D U-Net:

# train 2D U net
!nnUNet_train 2d nnUNetTrainerV2 1 0 --npz

训练级联模型:

# train 3D U-net cascade
!nnUNet_train 3d_lowres nnUNetTrainerV2CascadeFullRes 1 0 --npz
!nnUNet_train 3d_fullres nnUNetTrainerV2CascadeFullRes 1 0 --npz

例如:

#train 3D full resolution U net
!nnUNet_train 3d_fullres nnUNetTrainerV2 1 0 --npz 

推理

现在我们可以运行推论:

result_dir = os.path.join(task_folder_name, 'nnUNet_Prediction_Results')
make_if_dont_exist(result_dir, overwrite=True)
# -i is the input folder
# -o is where you want to save the predictions
# -t 1 means task 1, change it if you have a different task number
# Use -m 2d, or -m 3d_fullres, or -m 3d_cascade_fullres
!nnUNet_predict -i /content/drive/MyDrive/neurosciences-segmentation/my_nnunet/nnUNet/nnunet/nnUNet_raw_data_base/nnUNet_raw_data/Task001/imagesTs -o /content/drive/MyDrive/neurosciences-segmentation/my_nnunet/nnUNet/nnunet/nnUNet_raw_data_base/nnUNet_raw_data/Task001/nnUNet_Prediction_Results -t 1 -tr nnUNetTrainerV2 -m 2d -f 0  --num_threads_preprocessing 1

预测的可视化

首先让我们检查一下训练损失。这看起来很健康，我们有一个骰子得分Dice Score > 0.9(绿色曲线)。

这对于如此少的工作和3D神经成像分割任务来说确实是非常出色的。

我们来看一个示例:

很明显，这个模型已经有效地学会了如何高精度地分割大脑图像。虽然可能存在一些小的缺陷，但重要的是，图像分割领域正在迅速发展，我们正在朝着完美迈进。