[医学影像分割] nnunet 处理自己的数据集 + 并行训练 +训练时遇到的问题

【代码】[医学影像分割] nnunet 如何创建自己的数据集？

只是鸽子

738人浏览 · 2024-07-14 17:10:31

只是鸽子 · 2024-07-14 17:10:31 发布

version-1

使用这个脚本，来为你自己的数据创建dataset.json
生成一个Taskxxx的文件夹

import os
import shutil
import json

dataset_json = {
    "name" : "HNC Segmentation",
    "description" : "HNC pre-radiological treatment Segmentation",
    "tensorImageSize" : "3D",
    "modality": {
        "0": "T2"
    },
    "labels" : {
        "0": "background",
        "1": "GTVp",
        "2": "GTVn"
    },
    "file_ending": ".nii.gz",
    "numTraining" : 0,
    "numTest" : 0,
    "training" : [],
    "validation": [],
    "test" : []
}



results_path = "/mnt/home/pc/Chenq_team/HNC_dataset/nnunetv1_dataset/nnUNet_raw_data/Task018_HNC/"
dir_path = "/mnt/home/pc/Chenq_team/HNC_dataset/HNTSMRG24_train"

os.makedirs(results_path, exist_ok=True)
os.makedirs(os.path.join(results_path, "imagesTr"), exist_ok=True)
os.makedirs(os.path.join(results_path, "labelsTr"), exist_ok=True)
os.makedirs(os.path.join(results_path, "imagesTs"), exist_ok=True)


for dir in os.listdir(dir_path):
    pre_rt_dir_path = os.path.join(dir_path, dir, "preRT")
    pid = int(dir)
    img_new_path_json, mask_new_path = None, None
    for file in os.listdir(pre_rt_dir_path):
        if "T2" in file:
            img_old_path = os.path.join(pre_rt_dir_path, file)
            img_new_path = os.path.join(results_path, "imagesTr", f"preRT_{pid:03d}_0000.nii.gz")
            img_new_path_json = os.path.join(results_path, "imagesTr", f"preRT_{pid:03d}.nii.gz")
            shutil.copy(img_old_path, img_new_path)
        if "mask" in file:
            mask_old_path = os.path.join(pre_rt_dir_path, file)
            mask_new_path = os.path.join(results_path, "labelsTr", f"preRT_{pid:03d}.nii.gz")
            shutil.copy(mask_old_path, mask_new_path)
    assert img_new_path is not None and mask_new_path is not None, f"{pre_rt_dir_path} does not contain T2 or mask file"
    dataset_json["training"].append(
        {
            "image": img_new_path_json,
            "label": mask_new_path
        },
    )
    dataset_json["numTraining"] += 1

with open(os.path.join(results_path, "dataset.json"), "w") as f:
    json.dump(dataset_json, f, indent=4)

常见的命令

nnUNet_plan_and_preprocess -t 19 --verify_dataset_integrity  # 预处理数据集

自己划分数据集

更换split_final.pkl文件

version-1的并行训练

我没有找到相关的nnunet命令，但是nnunet-v1提供了run_training_DDP.py文件。训练命令如下:
其中MedNeXt没有提供DDP的trainer，自己抄一下nnunet的trainer就好了。记得替换sync_batchnorm，然后注意一下是半精度训练还是fp32。

export nnUNet_raw_data_base="/home/Guanjq/HNC_SEG_Data/nnunetv1_dataset/" # replace to your database
export nnUNet_preprocessed="/home/Guanjq/HNC_SEG_Data/nnunetv1_dataset/nnUNet_preprocessed/"
export RESULTS_FOLDER="/home/Guanjq/HNC_SEG_Data/nnunetv1_dataset/nnUNet_results/"

export TORCH_DISTRIBUTED_DEBUG="DETAIL"
export fold=3
export CUDA_VISIBLE_DEVICES=0,1 
export WORLD_SIZE=2
python3 -m torch.distributed.launch \
    --master_port 23331 \
    --nproc_per_node=2 \
    /home/Guanjq/Work/MedNeXt/nnunet_mednext/run/run_training_DDP.py \
    -network 3d_fullres \
    -network_trainer nnUNetTrainerV2_MedNeXt_S_kernel3_DDP \
    -task Task018_HNC_pre \
    -fold ${fold} \
    -p /home/Guanjq/HNC_SEG_Data/nnunetv1_dataset/nnUNet_preprocessed/Task018_HNC_pre/nnUNetPlansv2.1 \
    --fp32

version-2

使用这个脚本，来为你自己的数据创建dataset.json
生成一个Taskxxx的文件夹

import os
import shutil
import json

dataset_json = {
    "name" : "HNC Segmentation",
    "description" : "HNC pre-radiological treatment Segmentation",
    "tensorImageSize" : "3D",
    "channel_names": {
        "0": "T2"
    },
    "labels" : {
        "background" : 0,
        "GTVp" : 1,
        "GTVn" : 2
    },
    "file_ending": ".nii.gz",
    "numTraining" : 0,
    "numTest" : 0,
    "training" : [],
    "validation": [],
    "test" : []
}



results_path = "/mnt/home/pc/Chenq_team/HNC_dataset/nnunet_dataset/raw/Dataset018_HNC/"
dir_path = "/mnt/home/pc/Chenq_team/HNC_dataset/HNTSMRG24_train"

os.makedirs(results_path, exist_ok=True)
os.makedirs(os.path.join(results_path, "imagesTr"), exist_ok=True)
os.makedirs(os.path.join(results_path, "labelsTr"), exist_ok=True)

for dir in os.listdir(dir_path):
    pre_rt_dir_path = os.path.join(dir_path, dir, "preRT")
    pid = int(dir)
    img_new_path, mask_new_path = None, None
    for file in os.listdir(pre_rt_dir_path):
        if "T2" in file:
            img_old_path = os.path.join(pre_rt_dir_path, file)
            img_new_path = os.path.join(results_path, "imagesTr", f"preRT_{pid:03d}_0000.nii.gz")
            shutil.copy(img_old_path, img_new_path)
        if "mask" in file:
            mask_old_path = os.path.join(pre_rt_dir_path, file)
            mask_new_path = os.path.join(results_path, "labelsTr", f"preRT_{pid:03d}.nii.gz")
            shutil.copy(mask_old_path, mask_new_path)
    assert img_new_path is not None and mask_new_path is not None, f"{pre_rt_dir_path} does not contain T2 or mask file"
    dataset_json["training"].append(
        {
            "image": img_new_path,
            "label": mask_new_path
        },
    )
    dataset_json["numTraining"] += 1

with open(os.path.join(results_path, "dataset.json"), "w") as f:
    json.dump(dataset_json, f, indent=4)

遇到的问题

RuntimeError: One or more background workers are no longer alive. Exiting. Please check the print statements above for the actual error message

在使用version-1训练时，报了这个错误，尝试了 export nnUNet_n_proc_DA="1" 都没办法解决问题。

version_2 常用的命令

（1）nnUNetv2_convert_MSD_dataset -h

nnUNetv2_convert_MSD_dataset -h  # 查看这个命令的用法
nnUNetv2_convert_MSD_dataset -i /mnt/Liuyh/work/nnUNet/nnUNetData/nnUNet_raw/Task930_SixFeatures
# 这个命令会在Task文件夹同一个目录下面生成一个Dataset930文件夹，这个文件夹是后面用到的数据，他是初步检查了一下模态信息，也就是为每一个nii.gz文件添加了一个_0000或者_0001的一个模态后缀

（2）nnUNetv2_preprocess
这个命令的作用是初步处理所有的3D的nii.gz文件，为他们初始化处理了一下（主要是spacing的统一、以及像素值的归一化吧）

nnUNetv2_preprocess -d [数据集ID或名称] -pl [计划文件路径] -i [输入文件夹路径] -o [输出文件夹路径]  # 自己设置输入的文件夹路径，同时数据集ID必须如何这个路径，同时自己设置输出路径
或者
nnUNetv2_preprocess -d [数据集ID或名称] -pl [计划文件路径]  # 利用已有的plan文件
或者
nnUNetv2_preprocess -d [数据集ID或名称]  # 自动根据dataset.json生成plan文件

（3）train
可以设定训练的模型，训练的fold

CUDA_VISIBLE_DEVICES=2 nnUNetv2_train 950 3d_fullres 0

（4）推理

推理过程是一定要指定数据集ID、训练器类型、输入文件路径的，同时这个输入的文件一定要有模态后缀，跟Datasetxxx文件夹里面的数据的命名一样（主要是模态要对上）。

预测的时候，数据不需要经过preprocess，也就是说不一定放在imagesTs文件夹，可以是任意一个路径，只要image的名字合理即可。

nnUNetv2_predict -d [数据集ID或名称] -c [训练器类型]  -i [预处理后的输入文件夹路径] -o [输出文件夹路径] 
或者
nnUNetv2_predict -d [数据集ID或名称] -c [训练器类型]  -i [预处理后的输入文件夹路径] -o [输出文件夹路径] -f 0
# 额外设定使用fold0进行预测

比如我在淋巴结多标签分类数据集上就用了这些例子：

# 预测肌肉
nnUNetv2_predict -d 910 -c 3d_fullres -i /mnt/Liuyh/work/ALL_T2/zero_mask_image/ -o "/mnt/Liuyh/work/ALL_T2/943257output_muscle/"
# 预测肿瘤
nnUNetv2_predict -d 950 -c 3d_fullres -i /mnt/Liuyh/work/ALL_T2/zero_mask_image/ -o "/mnt/Liuyh/work/ALL_T2/943257output_muscle/"

如果预测出来什么都没有，很可能是图片本身就被截取了这部分区域

魔乐社区

魔乐社区（Modelers.cn) 是一个中立、公益的人工智能社区，提供人工智能工具、模型、数据的托管、展示与应用协同服务，为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作，由全产业链共同建设、共同运营、共同享有，推动国产AI生态繁荣发展。

更多推荐