nnUNetv2实战:从零构建二维医学影像分割数据集全流程

张开发
2026/4/12 9:01:45 15 分钟阅读

分享文章

nnUNetv2实战:从零构建二维医学影像分割数据集全流程
1. 环境准备搭建nnUNetv2开发环境第一次接触nnUNetv2时我花了整整两天时间折腾环境配置。现在回想起来其实只要抓住几个关键点就能避免90%的坑。建议使用Linux系统Ubuntu 20.04最佳Windows子系统也行但性能会打折扣。先说说Python版本的选择。官方推荐≥3.9但我实测3.10更稳定。创建环境时建议单独指定python版本conda create -n nnunet python3.10 -y conda activate nnunetPyTorch安装是第一个容易翻车的地方。很多教程只告诉你要装≥2.0版本但没强调CUDA版本匹配的重要性。先用nvidia-smi查显卡驱动支持的CUDA最高版本比如我实验室的RTX 3090显示CUDA 11.7那么就该这样安装pip install torch2.0.1cu117 torchvision0.15.2cu117 --extra-index-url https://download.pytorch.org/whl/cu117安装nnUNetv2本体时有个隐藏技巧先装依赖再克隆仓库。因为setup.py里有些依赖版本可能过时我整理了个增强版的依赖列表pip install tqdm numpy scipy pandas matplotlib seaborn scikit-learn scikit-image SimpleITK git clone https://github.com/MIC-DKFZ/nnUNet.git cd nnUNet pip install -e .验证安装是否成功时别只看import有没有报错。跑个简单测试更靠谱from nnunetv2.run.run_training import run_training print(导入成功)2. 数据准备构建符合规范的二维数据集去年处理乳腺超声数据集时我踩遍了数据格式的坑。nnUNetv2虽然支持二维图像了但文件结构和命名规则仍然严格。建议按这个目录树组织原始数据raw_data/ ├── train │ ├── images/ # 原图 │ └── masks/ # 标注图 └── val ├── images/ └── masks/关键来了所有图像必须统一格式。虽然官方说支持png/jpg但实测png最稳定。标注mask必须满足单通道灰度图像素值只能是0背景和255目标与原图严格同名如原图case1.png对应maskcase1.png用Python批量检查mask合规性from PIL import Image import numpy as np def validate_mask(mask_path): img Image.open(mask_path) arr np.array(img) unique_values np.unique(arr) assert set(unique_values).issubset({0,255}), f非法像素值{unique_values}3. 数据转换定制化适配脚本官方提供的转换脚本需要大改才能用。我以眼底血管分割为例分享个通用模板import os from batchgenerators.utilities.file_and_folder_operations import * from nnunetv2.dataset_conversion.generate_dataset_json import generate_dataset_json def convert_2d_dataset(source_dir, nnUNet_raw, dataset_nameDataset120_RetinaVessel): # 创建标准目录结构 imagestr join(nnUNet_raw, dataset_name, imagesTr) imagests join(nnUNet_raw, dataset_name, imagesTs) labelstr join(nnUNet_raw, dataset_name, labelsTr) maybe_mkdir_p(imagestr) maybe_mkdir_p(imagests) maybe_mkdir_p(labelstr) # 处理训练集 train_images subfiles(join(source_dir, train, images), suffix.png) for img_path in train_images: case_id os.path.basename(img_path)[:-4] # 移除.png mask_path join(source_dir, train, masks, f{case_id}.png) # 复制并重命名图像添加_0000后缀 target_img_path join(imagestr, f{case_id}_0000.png) copyfile(img_path, target_img_path) # 处理mask target_mask_path join(labelstr, f{case_id}.png) process_mask(mask_path, target_mask_path) # 你的mask处理函数 # 生成dataset.json generate_dataset_json( output_folderjoin(nnUNet_raw, dataset_name), channel_names{0: RGB}, labels{background: 0, vessel: 1}, num_training_caseslen(train_images), file_extension.png, dataset_namedataset_name )注意三个死亡陷阱数据集ID必须是Dataset开头三位数≥100图像文件名必须有_0000后缀即使单通道dataset.json里的channel_names要和实际数据匹配4. 预处理与训练二维任务优化技巧设置环境变量时建议直接写入.bashrc避免每次重设echo export nnUNet_raw你的raw路径 ~/.bashrc echo export nnUNet_preprocessed你的preprocessed路径 ~/.bashrc echo export nnUNet_results你的results路径 ~/.bashrc source ~/.bashrc运行预处理时加个--verbose参数能看到更多细节nnUNetv2_plan_and_preprocess -d 120 --verify_dataset_integrity --verbose二维训练的关键参数组合以五折交叉验证为例for fold in 0 1 2 3 4; do nnUNetv2_train 120 2d $fold --npz -num_epochs 300 done实测有效的训练技巧在nnUNet/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py里修改default_num_epochs小数据集(100样本)建议增大oversample_foreground_percent图像尺寸差异大时开启use_mask_for_norm5. 实战调试常见报错解决方案报错1Could not detect any training cases检查imagesTr和labelsTr的文件名对应关系确认dataset.json中的numTraining与实际数量一致报错2ValueError: Malformed segmentaion map用PIL重新保存maskImage.fromarray(mask).save(fixed.png)检查是否混入了RGBA格式报错3CUDA out of memory在nnUNet/nnunetv2/run/run_training.py中调整batch_size添加--disable_checkpointing节省显存我在乳腺肿瘤分割任务中发现适当调整这些参数能提升5%左右的Dice分数# 在Trainer类中添加 def __init__(self, plans, configuration, fold, dataset_json): super().__init__(plans, configuration, fold, dataset_json) self.batch_size 16 # 默认8 self.patch_size [256, 256] # 根据图像尺寸调整最后提醒nnUNetv2会自动选择最佳模型但中间结果保存在nnUNet_results/DATASET_ID/nnUNetTrainer__nnUNetPlans__2d/fold_X/checkpoints里遇到训练中断可以手动加载checkpoint继续训练。

更多文章