nnUNetv2训练二维图像数据集

目前nnUNetv2已经可以支持二维图像的生成，不用再进行复杂的二维图像转三维图像再进行训练，最近我使用nnUNet训练了自己的二维图像数据集，遇到了一定的困难，现将步骤总结方便大家使用，第一次分享，如有不正确的地方请大家及时指出~

nnunetv2：MIC-DKFZ/nnUNet (github.com)

一、虚拟环境创建以及环境配置

1.创建虚拟环境并激活虚拟环境

根据github上文档给出的建议，最好使用python>=3.9的版本。

#创建虚拟环境
conda create -n nnunet python=3.9
#激活虚拟环境
conda activate nnunet

2.安装pytorch

根据github上文档给出的建议，最好使用torch>=2.0的版本，结合自己gpu的cuda版本安装合适的pytorch，我这里使用的命令如下

pip install torch==2.0.0+cu117 torchvision==0.15.1+cu117 torchaudio==2.0.1 --index-url https://download.pytorch.org/whl/cu117

3.安装nnUNet

首先在当前目录下生成nnUNet的仓库

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

进入nnUNet文件夹

cd nnUNet

安装所需依赖包

pip install -e .

至此，环境配置就顺利完成了

二、数据转换

1.导入自己的数据集

将自己的数据集文件夹放在nnUNet目录下，并且提前自行分好训练集和测试集（我这里分别写成的是train和val）

2.修改数据转换文件，得到nnUNetv2适配的数据集格式

打开Dataset120_RoadSegmentation.py文件

 

根据自己数据集的格式修改代码，我的原始数据集中图像名称如下：

 我修改后的主函数代码如下：

if __name__ == "__main__":
    # extracted archive from https://www.kaggle.com/datasets/insaff/massachusetts-roads-dataset?resource=download
    #自己数据集存放路径
    source = '/aidata/YinHanLong/nnUNet/data_train_valid'
    #nnUNet_raw的路径
    nnUNet_raw = '/aidata/YinHanLong/nnUNet/nnUNet_raw'
    #自定义dataset_name的名称
    dataset_name = 'Dataset110_breastTumorSegmentation'

    imagestr = join(nnUNet_raw, dataset_name, 'imagesTr')
    imagests = join(nnUNet_raw, dataset_name, 'imagesTs')
    labelstr = join(nnUNet_raw, dataset_name, 'labelsTr')
    labelsts = join(nnUNet_raw, dataset_name, 'labelsTs')
    maybe_mkdir_p(imagestr)
    maybe_mkdir_p(imagests)
    maybe_mkdir_p(labelstr)
    maybe_mkdir_p(labelsts)
    
    #训练集和测试集的路径
    train_source = join(source, 'train')
    test_source = join(source, 'val')

    with multiprocessing.get_context("spawn").Pool(8) as p:

        # not all training images have a segmentation
        valid_ids = subfiles(join(train_source, 'mask'), join=False, suffix='png')
        num_train = len(valid_ids)
        r = []
        #导入训练集图像和label
        for v in valid_ids:
            r.append(
                p.starmap_async(
                    load_and_covnert_case,
                    ((
                         join(train_source, 'images', v.replace('_mask', '')),
                         join(train_source, 'mask', v),
                         join(imagestr, v[:-9] + '_0000.png'),
                         join(labelstr, v),
                         50
                     ),)
                )
            )

        # test set
        valid_ids = subfiles(join(test_source, 'mask'), join=False, suffix='png')
        #导入测试集图像和label
        for v in valid_ids:
            r.append(
                p.starmap_async(
                    load_and_covnert_case,
                    ((
                         join(test_source, 'images', v.replace('_mask', '')),
                         join(test_source, 'mask', v),
                         join(imagests, v[:-9] + '_0000.png'),
                         join(labelsts, v),
                         50
                     ),)
                )
            )
        _ = [i.get() for i in r]

    generate_dataset_json(join(nnUNet_raw, dataset_name), {0: 'R', 1: 'G', 2: 'B'}, {'background': 0, 'tumor': 1},
                          num_train, '.png', dataset_name=dataset_name)

3.运行程序：运行Dataset120_RoadSegmentation.py文件，在nnUNet_raw中得到转换过后的数据集

 修改过程中，需要注意以下几点：

（1）dataset_name的名称必须按照“Task”+数字+“_”+"名称"的格式命名，例如我的就是'Dataset110_breastTumorSegmentation'，这个数字需要大于10，并且牢记这个id。

（2）在数据集转换之后需要注意，label文件夹中文件名称必须与image文件夹中文件名称删除“_0000”之后的一致。如下：

（3）生成的dataset.json文件如下：

（4）注意自己的mask的格式

mask中像素灰度值必须只有0和255两个值，最近在尝试另外一个数据集的时候发现了这一点嘿嘿

三、数据预处理

1.设置环境变量

根据自己创建的文件夹位置，设置以下三个环境变量

#设置环境变量
export nnUNet_raw="/aidata/YinHanLong/nnUNet/nnUNet_raw"
export nnUNet_preprocessed="/aidata/YinHanLong/nnUNet/nnUNet_preprocessed"
export nnUNet_results="/aidata/YinHanLong/nnUNet/nnUnet_results"

 2.运行预处理命令

这里的DATASET_ID就是数据转换步骤中你设置的datasetname中的id。

nnUNetv2_plan_and_preprocess -d DATASET_ID --verify_dataset_integrity

 这一步如果报错，一定检查你的数据集下文件名称，一定要严格按照我上述步骤中说的方式命名。

 四、训练

nnUNetv2_train DATASET_NAME_OR_ID UNET_CONFIGURATION FOLD

 DATASET_NAME_OR_ID这里填写你的id即可，二维图像的话UNET_CONFIGURATION就填2d，FOLD代表几折交叉验证，我选择的是五折交叉验证，这里写5就好。

至此，使用nnUNetv2运行自己的二维数据集步骤就完成了，欢迎大家提问讨论~