目前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运行自己的二维数据集步骤就完成了,欢迎大家提问讨论~