赞
踩
本文是对nnUNet如何在2D数据上进行训练以及预测,是对几篇博文参考之后的总结,由于本人对nnUNet使用也是刚开始不久,所以对于nnUNet的理解肯定还停留在表层,所以原理上的东西就不多介绍了,更侧重于介绍其在2D数据上的使用流程。官方也有nnUNet训练2D数据的例子,一切以官方为准,本文仅作参考。
本人的系统环境为:Ubuntu18.04,pytorch为:1.7.1+cu110。其余环境没有测试过。
参考文章:
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大小。
由于本人避免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开始工作的指令。
接下来是创建数据保存目录。
假定我们的数据长这样,图片和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:
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。
首先需要将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
运行完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生成五折交叉验证的结果。
因为模型预测的结果时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)
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。