开发: C++知识库 Java知识库 JavaScript Python PHP知识库人工智能区块链大数据移动开发嵌入式开发工具数据结构与算法开发测试游戏开发网络协议系统运维
教程: HTML教程 CSS教程 JavaScript教程 Go语言教程 JQuery教程 VUE教程 VUE3教程 Bootstrap教程 SQL数据库教程 C语言教程 C++教程 Java教程 Python教程 Python3教程 C#教程
数码: 电脑笔记本显卡显示器固态硬盘硬盘耳机手机 iphone vivo oppo 小米华为单反装机图拉丁

-> 人工智能 -> nnUNet在2D图像上的训练和测试的使用教程 -> 正文阅读

[人工智能]nnUNet在2D图像上的训练和测试的使用教程

写在前面的话

本文是对nnUNet如何在2D数据上进行训练以及预测，是对几篇博文参考之后的总结，由于本人对nnUNet使用也是刚开始不久，所以对于nnUNet的理解肯定还停留在表层，所以原理上的东西就不多介绍了，更侧重于介绍其在2D数据上的使用流程。官方也有nnUNet训练2D数据的例子，一切以官方为准，本文仅作参考。
本人的系统环境为：Ubuntu18.04，pytorch为：1.7.1+cu110。其余环境没有测试过。
参考文章：

一、 nnUNet简介

nnUNet是一种基于2D and 3D 原始U-Nets 自适应框架，该框架能根据给定数据集的属性自动调整所有超参数，整个过程无需人工干预。仅仅依赖于朴素的U-Net结构（就是原始U-Net）和鲁棒的训练方案，nnU-Net在六个得到公认的分割挑战中实现了最先进的性能。
github仓库：https://github.com/MIC-DKFZ/nnUNet
paper：《Automated Design of Deep Learning Methods for Biomedical Image Segmentation》

nnUNet是德国癌症研究中心的工程师编写的框架，迄今为止依旧在维护和更新。虽然整个过程实现自动化，但是对GPU资源要求比较高，虽然官方说For training nnU-Net models the GPU should have at least 10 GB，但是实测后发现，如果按默认配置，需要19G显存。如果显存不足可以考虑降低batch size大小。

二、nnUNet安装

由于本人避免python环境冲突，用conda创建一个虚拟环境nnUNet，并激活虚拟环境

conda create -n nnUNet python=3.8
source activate nnUnet

由于是面向初学者的使用教程，初学者请务必按照我的做法，等你熟练掌握以后再考虑新的姿势（有些文件夹的创建时多余的，但是你还是跟着我这样做最好）。
第一步：在某个目录下(比如说/data)创建nnUNetFrame目录，在nnUNet的终端环境下cd到这个nnUNetFrame目录，然后git clone https://github.com/MIC-DKFZ/nnUNet.git
第二步：cd nnUNet
第三步：pip install -e .（兄弟们和集美们别忘了加 . ）
当你安装完成这些以后，你的每一次对nnUNet的操作，都会在命令行里以nnUNet_开头，代表着你的nnUNet开始工作的指令。

接下来是创建数据保存目录。

进入你之前创建的nnUNetFrame文件夹里面，创建一个名为DATASET的文件夹，现在你的nnUNetFrame文件夹下有两个文件夹，nnUNet是代码源，另一个DATASET就是我们接下来用来放数据的地方；
进入创建好的DATASET文件夹下面，创建下面三个文件夹分别是nnUNet_preprocessed、nnUNet_raw、nnUNet_trained_models

第一个用来存放原始数据预处理之后的数据，第二个用来存放原始的你要训练的数据，第三个用来存放训练的结果。
进入上面第二个文件夹nnUNet_raw，创建下面两个文件夹，分别是nnUNet_cropped_data、nnUNet_cropped_data，右边为原始数据，左边为crop以后的数据。
进入右边文件夹nnUNet_raw_data，创建一个名为Task001_BloodVessel的文件夹（解释：这nnUnet的数据格式是固定的，Task001_BloodVessel由Task+ID+数据名组成，你可以对这个任务的数字ID进行任意的命名，比如你要分割心脏，你可以起名为Task001_Heart，比如你要分割肾脏，你可以起名为Task002_Kidney，前提是必须按照这种格式）

三、数据准备

3.1 2D数据转3D数据

假定我们的数据长这样，图片和mask分别存放在img和mask的文件夹下，两者同名。
在这里插入图片描述
由于这是2D图片数据，要将其转为3D数据，其实就是z轴为1的3维数据。
所以随意新建一个项目，创建一个2DDataProcessTo3D.py脚本

import os
import random
from tqdm import tqdm
import SimpleITK as sitk
import cv2
import numpy as np


def SplitDataset(img_path, train_percent=0.9):
    data = os.listdir(img_path)
    train_images = []
    test_images = []
    num = len(data)
    train_num = int(num * train_percent)
    indexes = list(range(num))
    train = random.sample(indexes, train_num)
    for i in indexes:
        if i in train:
            train_images.append(data[i])
        else:
            test_images.append(data[i])
    return train_images, test_images


def conver(img_path,  save_dir, mask_path=None, select_condition=None, mode="trian"):
    os.makedirs(save_dir, exist_ok=True)
    if mode == "train":
        savepath_img = os.path.join(save_dir, 'imagesTr')
        savepath_mask = os.path.join(save_dir, 'labelsTr')
    elif mode == "test":
        savepath_img = os.path.join(save_dir, 'imagesTs')
        savepath_mask = os.path.join(save_dir, 'labelsTs')
    os.makedirs(savepath_img, exist_ok=True)
    if mask_path is not None:
        os.makedirs(savepath_mask, exist_ok=True)

    ImgList = os.listdir(img_path)
    with tqdm(ImgList, desc="conver") as pbar:
        for name in pbar:
            if select_condition is not None and name not in select_condition:
                continue
            Img = cv2.imread(os.path.join(img_path, name))
            if mask_path is not None:
                Mask = cv2.imread(os.path.join(mask_path, name), 0)
                Mask = (Mask / 255).astype(np.uint8)
                if Img.shape[:2] != Mask.shape:
                    Mask = cv2.resize(Mask, (Img.shape[1], Img.shape[0]))
            Img_Transposed = np.transpose(Img, (2, 0, 1))
            Img_0 = Img_Transposed[0].reshape(1, Img_Transposed[0].shape[0], Img_Transposed[0].shape[1])
            Img_1 = Img_Transposed[1].reshape(1, Img_Transposed[1].shape[0], Img_Transposed[1].shape[1])
            Img_2 = Img_Transposed[2].reshape(1, Img_Transposed[2].shape[0], Img_Transposed[2].shape[1])
            if mask_path is not None:
                Mask = Mask.reshape(1, Mask.shape[0], Mask.shape[1])
    
            Img_0_name = name.split('.')[0] + '_0000.nii.gz'
            Img_1_name = name.split('.')[0] + '_0001.nii.gz'
            Img_2_name = name.split('.')[0] + '_0002.nii.gz'
            if mask_path is not None:
                Mask_name = name.split('.')[0] + '.nii.gz'
    
            Img_0_nii = sitk.GetImageFromArray(Img_0)
            Img_1_nii = sitk.GetImageFromArray(Img_1)
            Img_2_nii = sitk.GetImageFromArray(Img_2)
            if mask_path is not None:
                Mask_nii = sitk.GetImageFromArray(Mask)
    
            sitk.WriteImage(Img_0_nii, os.path.join(savepath_img, Img_0_name))
            sitk.WriteImage(Img_1_nii, os.path.join(savepath_img, Img_1_name))
            sitk.WriteImage(Img_2_nii, os.path.join(savepath_img, Img_2_name))
            if mask_path is not None:
                sitk.WriteImage(Mask_nii, os.path.join(savepath_mask, Mask_name))


if __name__ == "__main__":
    train_percent = 0.9
    img_path = r".../img"
    mask_path = r".../mask"
    output_folder = r"./dataset"
    os.makedirs(output_folder, exist_ok=True)
    train_images, test_images = SplitDataset(img_path, train_percent)
    conver(img_path, output_folder, mask_path, train_images, mode="train")
    conver(img_path, output_folder, mask_path, test_images, mode="test")

一些参数解释：

train_percent：训练集和验证集划分比例
img_path：图片路径
mask_path：mask路径
output_folder：保存路径

几个注意的点：
imagesTr是训练数据，imagesTs是测试数据，labelsTr是训练数据的标签，labelsTs是测试数据的标签（这个可能没什么用），数据样本la_003_0000.nii.gz由case样本名字+模态标志+.nii.gz组成，不同的模态用0000/0001/0002/0003区分。
示例树结构：

nnUNet_raw_data_base/nnUNet_raw_data/Task001_BloodVessel
├── dataset.json
├── imagesTr
│   ├── la_003_0000.nii.gz
│   ├── la_004_0000.nii.gz
│   ├── ...
├── imagesTs
│   ├── la_001_0000.nii.gz
│   ├── la_002_0000.nii.gz
│   ├── ...
└── labelsTr
    ├── la_003.nii.gz
    ├── la_004.nii.gz
    ├── ...

我们的原始2维数据是RGB三通道的，我们可以把RGB三通道的数据看成3个模态，分别提取不同通道的数据，形状转换成（1，width, height），采用SimpleITK保存为3维数据。
imagesTr：
在这里插入图片描述
labelsTr：

3.2 制作dataset.json

dataset.json的文件内容是包含数据集的元数据，如任务名字，模态，标签含义，训练集包含的图像地址。
这里讲一下，dataset.json这个文件怎么弄。一个合格的文件应该包含如下信息：

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

在这里插入图片描述
最简单的方法就是copy别人的代码，在前人的基础上修改一下。如： dataset_conversion
最终效果如图（图片来源：保姆级教程：nnUnet在2维图像的训练和测试）：

dataset.json里的数据名去掉了模态标志，所以每个数据看起来有三个重复名字，在后面读图的时候会自动添加模态标志。还有name这一行是你任务数据集的名字，比如说，你在nnUNet_raw_data创建一个名为Task001_BloodVessel，那么这个name就是BloodVessel。

3.3 设置数据集环境变量

首先需要将2DDataProcessTo3D.py脚本中output_folder目录下的，imagesTr、imagesTs、labelsTr、labelsTs和制作好的dataset.json复制到nnUNet_raw_data/Task001_BloodVessel（假定你任务的名字就是这个，当然还可以命名成其他，具体请参照上面）。
在终端中运行导出命令，设置环境变量：

export nnUNet_raw_data_base="/data/nnUnet/Data/nnUNet_raw"
export nnUNet_preprocessed="/data/nnUnet/Data/nnUNet_preprocessed"
export RESULTS_FOLDER="/data/nnUnet/Data/nnUNet_trained_models"

假定你的目录是创建在/data下的。

四、数据预处理

nnUNet_plan_and_preprocess -t 100 -data23d 2 --verify_dataset_integrity

这里根据保姆级教程：nnUnet在2维图像的训练和测试所说增加了一个设置参数data23d用来指定要处理的是2维数据还是3维数据，在nnunet/preprocessing/sanity_checks.py中的verify_dataset_integrity函数中增加做如下设置：

 if data23d == '2':
     expected_train_identifiers = np.unique(expected_train_identifiers)
     expected_test_identifiers = np.unique(expected_test_identifiers)
     print('train num', len(expected_train_identifiers))
     print('test num:', len(expected_test_identifiers))

具体位置为：
在这里插入图片描述
此外，由于新增了一个参数，还需要在nnunet/experiment_planning/nnUNet_plan_and_preprocess.py中添加

efault=3, help="the data is 2d or 3d")

并且在传参一行要加入

if args.verify_dataset_integrity:
            verify_dataset_integrity(join(nnUNet_raw_data, task_name), args.data23d)

五、训练

训练时逐条运行下面命令

CUDA_VISIBLE_DEVICES=1 nnUNet_train 2d nnUNetTrainerV2 Task001_BloodVessel 0  --npz
CUDA_VISIBLE_DEVICES=1 nnUNet_train 2d nnUNetTrainerV2 Task001_BloodVessel 1  --npz
CUDA_VISIBLE_DEVICES=1 nnUNet_train 2d nnUNetTrainerV2 Task001_BloodVessel 2  --npz
CUDA_VISIBLE_DEVICES=1 nnUNet_train 2d nnUNetTrainerV2 Task001_BloodVessel 3  --npz
CUDA_VISIBLE_DEVICES=1 nnUNet_train 2d nnUNetTrainerV2 Task001_BloodVessel 4  --npz

六、预测

6.1 模型预测

运行完5折交叉验证后，可以确定最佳的配置，下面的008是Task的ID，我们这个任务是008

nnUNet_find_best_configuration -m 2d -t 008 –strict

运行后会在路径nnUNet_trained_models/nnUNet/ensembles/Task001_BloodVessel下生成
在这里插入图片描述
然后打开上面的txt文件，里面会生成预测方法：

nnUNet_predict -i FOLDER_WITH_TEST_CASES -o OUTPUT_FOLDER_MODEL1 -tr nnUNetTrainerV2 -ctr nnUNetTrainerV2CascadeFullRes -m 2d -p nnUNetPlansv2.1 -t Task001_BloodVessel

INPUT_FOLDER: 测试数据地址
OUTPUT_FOLDER：分割数据存放地址
CONFIGURATION：使用的什么架构，2d or 3d_fullres or 3d_cascade_fullres等

所以最终的预测命令就是：

nnUNet_predict -i /data/nnUnet/Data/nnUNet_raw/nnUNet_raw_data/Task001_BloodVessel/imagesTs/ -o /data/nnUnet/Data/nnUNet_raw/nnUNet_raw_data/Task001_BloodVessel/imagesTsPred/ -tr nnUNetTrainerV2 -ctr nnUNetTrainerV2CascadeFullRes -m 2d -p nnUNetPlansv2.1 -t Task08_HepaticVessel

就会在imagesTsPred生成五折交叉验证的结果。

6.2 还原结果

因为模型预测的结果时nii.gz文件，所以需要将其转成2D的图像数据
所以新建一个3DDataProcessTo2D.py的脚本，并计算dice

import os
from tqdm import tqdm
import SimpleITK as sitk
import cv2
import numpy as np
from sklearn.metrics import f1_score


def conver(img_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    img_list = [i for i in os.listdir(img_dir) if ".nii.gz" in i]
    with tqdm(img_list, desc="conver") as pbar:
        for name in pbar:
            image = sitk.ReadImage(os.path.join(img_dir, name))
            image = sitk.GetArrayFromImage(image)[0]
            image = (image * 255).astype(np.uint8)
            cv2.imwrite(os.path.join(output_dir, name.split(".")[0]+".png"), image)


def computer_dice(mask_dir, label_dir):
    dice_list = []
    name_list = os.listdir(mask_dir)
    with tqdm(name_list, desc="dice") as pbar:
        for name in pbar:
            mask = cv2.imread(os.path.join(mask_dir, name), cv2.IMREAD_GRAYSCALE)
            label = cv2.imread(os.path.join(label_dir, name), cv2.IMREAD_GRAYSCALE)
            if mask.shape != label.shape:
                mask = cv2.resize(mask, (label.shape[1], label.shape[0]))
            mask = (mask / 255).ravel().astype(np.int)
            label = (label / 255).ravel().astype(np.int)
            dice = f1_score(y_true=label, y_pred=mask)
            dice_list.append(dice)
    print(sum(dice_list)*1.0/len(dice_list))


if __name__ == "__main__":
    img_dir = "/data/nnUNetFrame/DATASET/nnUNet_raw/nnUNet_raw_data/Task001_BloodVessel/imagesTsPred"
    output_dir = "/data/nnUNetFrame/DATASET/nnUNet_raw/nnUNet_raw_data/Task001_BloodVessel/imagesTsPredMask"
    conver(img_dir, output_dir)
  
    label_dir = "..../train_dataset/mask"
    computer_dice(output_dir, label_dir)