nnUNetV2实战：从零开始处理医学影像2D数据集（附完整代码）

nnUNetV2实战从零构建医学影像2D分割全流程如果你刚接触医学影像分割面对五花八门的模型和复杂的预处理流程感到无从下手那么这篇文章就是为你准备的。nnUNetV2 不是一个新模型而是一套高度自动化、标准化的医学影像分割框架它把数据预处理、实验规划、模型训练和推理这些繁琐的步骤封装成了一条流水线。对于研究者来说这意味着你可以把更多精力放在数据质量和业务逻辑上而不是反复调试模型参数和数据处理脚本。这篇文章将带你完整走一遍使用 nnUNetV2 处理自定义 2D 医学影像数据集的全过程。我们会从最基础的环境搭建开始一步步解决数据格式转换、RGB标签处理、自定义数据集划分等实际工程中必然会遇到的“坑”。无论你是需要快速验证一个想法的算法工程师还是希望将深度学习技术应用于具体临床场景的研究人员这套流程都能为你提供一个坚实可靠的起点。1. 环境搭建与核心概念理解在开始敲代码之前我们需要先理清几个关键概念这能帮你更好地理解后续每一步操作的意义而不是机械地复制命令。nnUNetV2 的三个核心目录是整个框架运行的基石你必须正确设置它们nnUNet_raw: 这是你的原始数据仓库。所有未经 nnUNet 处理过的数据集都必须按照它规定的格式放在这里。你可以把它想象成原料库。nnUNet_preprocessed: 这是预处理中心。框架会自动读取nnUNet_raw中的数据进行归一化、重采样等操作并将处理后的中间文件存储于此。这是加工车间。nnUNet_results: 这是成品仓库。所有训练好的模型权重、训练日志、预测结果以及模型集成方案都会保存在这里。这是最终的成品库和档案室。这三个路径需要通过环境变量告诉 nnUNet它才知道去哪里找东西。接下来我们开始具体的环境配置。1.1 创建隔离的Python环境我强烈建议使用 Conda 来管理环境它能避免不同项目间的依赖冲突。这里我们使用 Python 3.10这是一个在稳定性和新特性之间取得较好平衡的版本。conda create -n nnunetv2 python3.10 -y conda activate nnunetv21.2 安装 PyTorch 与 nnUNetV2PyTorch 的安装需要根据你的 CUDA 版本进行选择。你可以通过nvidia-smi命令查看驱动支持的 CUDA 最高版本。以下命令适用于 CUDA 11.7pip install torch2.0.0cu117 torchvision0.15.1cu117 torchaudio2.0.1 --index-url https://download.pytorch.org/whl/cu117注意请务必访问 PyTorch 官方安装指南选择与你系统匹配的版本命令。安装错误版本的 PyTorch 是后续一切问题的根源。安装好 PyTorch 后安装 nnUNetV2 就非常简单了pip install nnunetv2这个命令会安装 nnUNetV2 的核心库以及一系列命令行工具。安装成功后你可以在终端尝试输入nnUNetv2_然后按 Tab 键如果能看到一系列以nnUNetv2_开头的命令补全说明安装成功。这些命令就是我们后续操作的主要工具。1.3 可选安装网络结构可视化工具HiddenLayer 是一个轻量级的工具可以让 nnUNet 在训练开始时绘制出网络结构图。这对于理解模型架构很有帮助但不是必须的。pip install --upgrade githttps://github.com/FabianIsensee/hiddenlayer.gitmore_plotted_details如果安装失败通常是因为缺少 Graphviz。你可以使用系统包管理器安装如sudo apt-get install graphviz或者直接跳过这一步不影响核心功能。1.4 配置环境变量这是至关重要的一步。你需要将前面提到的三个核心目录的绝对路径设置为环境变量。假设你的项目结构如下/home/yourname/medical_ai/ ├── nnUNet_raw/ ├── nnUNet_preprocessed/ └── nnUNet_results/那么你需要将以下内容添加到你的 shell 配置文件如~/.bashrc或~/.zshrc的末尾export nnUNet_raw/home/yourname/medical_ai/nnUNet_raw export nnUNet_preprocessed/home/yourname/medical_ai/nnUNet_preprocessed export nnUNet_results/home/yourname/medical_ai/nnUNet_results保存后执行source ~/.bashrc使配置生效。你可以通过echo $nnUNet_raw来验证变量是否设置正确。2. 准备你的自定义数据集从混乱到规范nnUNet 的强大之处在于其“约定大于配置”的理念但前提是你的数据必须遵循它的约定。对于 2D 图像如病理切片、视网膜图像、皮肤镜图像最常见的坑就是标签格式。2.1 理解 nnUNet 的原始数据格式在nnUNet_raw目录下每个数据集都必须放在一个以DatasetXXX_YYY命名的文件夹中其中XXX是大于 10 的唯一数字 IDYYY是数据集名称。例如Dataset120_MyLungCT。在这个数据集文件夹内结构必须是这样的Dataset120_MyLungCT/ ├── dataset.json ├── imagesTr/ │ ├── case_001_0000.png │ ├── case_002_0000.png │ └── ... └── labelsTr/ ├── case_001.png ├── case_002.png └── ...关键点解析imagesTr/: 存放训练图像。注意文件名格式{案例ID}_{模态ID}.{后缀}。对于 2D 单模态数据如普通照片模态 ID 固定为0000。所以你会看到case_001_0000.png。labelsTr/: 存放对应的训练标签即分割掩码。文件名只需{案例ID}.{后缀}与图像去掉_0000的部分对应。dataset.json: 数据集的元信息配置文件这是灵魂所在。2.2 处理棘手的 RGB 标签问题很多公开数据集或标注工具如 LabelMe产生的标签是 RGB 三通道的 PNG 文件。nnUNet 在内部处理时默认标签是单通道的每个像素值代表一个类别如 0-背景1-肿瘤。直接使用 RGB 标签会导致错误因为 nnUNet 会尝试读取三个通道发现数值不一致而报错。解决方法是在数据转换阶段将 RGB 标签转换为单通道的索引图。假设你的原始数据是这样一个混乱的文件夹my_raw_data/ ├── train_images/ │ ├── img1.jpg │ └── img2.jpg └── train_masks_rgb/ # 这里是RGB格式的PNG ├── img1.png └── img2.png你需要编写一个转换脚本。这里提供一个 Python 示例将 RGB 掩码假设白色[255,255,255]是前景黑色[0,0,0]是背景转换为单通道二值掩码import os from PIL import Image import numpy as np raw_image_dir “/path/to/my_raw_data/train_images” raw_mask_dir “/path/to/my_raw_data/train_masks_rgb” output_image_dir “/path/to/nnUNet_raw/Dataset120_MyLungCT/imagesTr” output_label_dir “/path/to/nnUNet_raw/Dataset120_MyLungCT/labelsTr” os.makedirs(output_image_dir, exist_okTrue) os.makedirs(output_label_dir, exist_okTrue) for img_name in os.listdir(raw_image_dir): case_id img_name.split(‘.’)[0] # 假设 img1.jpg - case_id‘img1’ # 1. 处理图像复制或转换到目标位置并重命名 img_src os.path.join(raw_image_dir, img_name) img_dst os.path.join(output_image_dir, f“{case_id}_0000.png”) # 如果是JPG可能需要用PIL打开再保存为PNG Image.open(img_src).save(img_dst) # 2. 处理RGB标签转换为单通道 mask_src os.path.join(raw_mask_dir, f“{case_id}.png”) # 假设同名 mask_dst os.path.join(output_label_dir, f“{case_id}.png”) mask_rgb np.array(Image.open(mask_src)) # 将白色像素(255,255,255)转为1其他转为0 # 这里逻辑根据你的RGB标签定义修改 mask_single_channel (mask_rgb[:, :, 0] 255).astype(np.uint8) Image.fromarray(mask_single_channel).save(mask_dst)2.3 编写 dataset.json 文件这个文件告诉 nnUNet 关于数据集的一切。一个典型的 2D 二分类数据集的dataset.json如下{ “channel_names”: { “0”: “RGB” }, “labels”: { “background”: 0, “tumor”: 1 }, “numTraining”: 100, “file_ending”: “.png”, “overwrite_image_reader_writer”: “PILWriter”, “image_reader_writer”: “PILWriter”, “modality”: { “0”: “RGB” }, “name”: “MyLungCT”, “reference”: “My own dataset”, “release”: “1.0”, “tensorImageSize”: “2D”, “description”: “2D lung CT slices for tumor segmentation” }参数解释表参数值示例说明channel_names{“0”: “RGB”}通道名称。对于2D RGB图像只有一个通道通道0但其内容是三通道的RGB信息。labels{“background”:0, “tumor”:1}关键定义类别名称与标签像素值的映射关系。背景必须是0。numTraining100训练案例的数量必须与imagesTr/下的文件数匹配。file_ending“.png”图像文件的后缀名。overwrite_image_reader_writer“PILWriter”对于2D图像PNG/JPG必须指定告诉nnUNet使用PIL库读写图像。tensorImageSize“2D”指明这是2D数据集。modality{“0”: “RGB”}描述模态。对于自然图像通常是RGB。特别注意overwrite_image_reader_writer和image_reader_writer设置为”PILWriter”是处理2D标准图像格式非医学专用格式如NIfTI的关键很多初学者在这里出错。3. 数据预处理与实验规划当你的Dataset120_MyLungCT文件夹在nnUNet_raw中准备就绪后就可以启动 nnUNet 的自动化预处理流水线了。3.1 执行规划与预处理只需一行命令nnUNet 就会分析你的数据并做好一切训练准备nnUNetv2_plan_and_preprocess -d 120 --verify_dataset_integrity-d 120: 指定你的数据集 ID即Dataset120中的 120。--verify_dataset_integrity: 让 nnUNet 先检查你的数据集格式和dataset.json是否正确。强烈建议首次运行时加上此参数。这个命令会执行以下操作数据完整性验证检查图像和标签是否一一对应标签值是否在dataset.json定义的范围内。指纹提取分析数据的像素间距、形状、强度分布等统计信息。实验规划基于“指纹”自动决定网络架构的配置如 patch size、batch size、训练轮数等。规划文件会保存在nnUNet_preprocessed/Dataset120_MyLungCT下。预处理根据规划对图像进行归一化、重采样如果需要等操作并将处理后的数据保存到nnUNet_preprocessed目录。这个过程可能会花费一些时间取决于数据量的大小。执行成功后你就可以在nnUNet_preprocessed目录下看到对应的文件夹和生成的nnUNetPlans.json文件。这个文件包含了为你数据集“量身定制”的训练方案通常不需要手动修改。4. 模型训练与交叉验证nnUNet 默认使用 5 折交叉验证来稳健地评估模型性能。训练命令非常直观nnUNetv2_train 120 2d 0120: 数据集 ID。2d: 使用 2D U-Net 配置。对于2D图像就选这个。0: 训练第 0 折。你需要分别运行 0, 1, 2, 3, 4 来训练完整的 5 折模型。训练开始后你会在终端看到实时日志包括损失、学习率、验证集 Dice 分数等。模型权重和日志默认保存在nnUNet_results/Dataset120_MyLungCT/nnUNetTrainer__nnUNetPlans__2d目录下每一折一个子文件夹fold_0,fold_1等。训练过程中的几个实用技巧监控训练除了看终端日志你可以实时查看nnUNet_results下对应 fold 文件夹里的progress.png它绘制了训练和验证损失曲线。恢复训练如果训练意外中断重新执行相同的命令nnUNet 会自动从最新的 checkpoint 恢复。资源调整如果遇到 GPU 内存不足nnUNet 自动规划的 patch size 可能过大。你可以手动修改nnUNet_preprocessed/Dataset120_MyLungCT/nnUNetPlans.json文件中的patch_size但需谨慎最好在规划阶段通过设置–optimizer_parameters来调整。5. 模型推理与后处理训练完 5 折模型后我们通常需要找到最佳配置并进行推理。5.1 寻找最佳配置与模型集成运行以下命令让 nnUNet 自动评估各折模型在验证集上的表现并给出集成建议nnUNetv2_find_best_configuration 120 -c 2d这个命令会输出一个详细的报告告诉你哪个 fold 的模型性能最好以及是否模型集成ensemble能带来提升。命令的最后它会直接给出推荐的最佳推理命令非常贴心。输出可能类似于***Best*: nnUNetTrainer__nnUNetPlans__2d: 0.912 ***Run inference like this:*** nnUNetv2_predict -d 120 -i INPUT_FOLDER -o OUTPUT_FOLDER -f 0 1 2 3 4 -tr nnUNetTrainer -c 2d -p nnUNetPlans这里-f 0 1 2 3 4表示使用全部 5 个 fold 的模型进行集成预测这通常能获得最鲁棒的结果。5.2 执行预测根据上面的输出你需要准备一个INPUT_FOLDER里面存放你需要预测的图像。这些图像也需要遵循同样的命名规范即{case_id}_0000.png并放在一个文件夹里例如nnUNet_raw/Dataset120_MyLungCT/imagesTs/。然后运行预测命令nnUNetv2_predict -d 120 \ -i /path/to/your/test_images/ \ -o /path/to/prediction_output/ \ -f 0 1 2 3 4 \ -tr nnUNetTrainer \ -c 2d \ -p nnUNetPlans预测完成后结果会以单通道 PNG 格式保存在输出目录像素值对应dataset.json中定义的标签。5.3 应用后处理nnUNet 在find_best_configuration阶段可能会学习到一个后处理模型例如移除小面积的孤立预测区域。如果推荐了后处理它会同时给出后处理命令。执行后处理可以进一步提升结果nnUNetv2_apply_postprocessing \ -i /path/to/prediction_output/ \ -o /path/to/prediction_output_pp/ \ -pp_pkl_file /path/to/postprocessing.pkl \ -np 8 \ -plans_json /path/to/plans.json其中-pp_pkl_file和-plans_json的路径可以在nnUNet_results下对应模型的crossval_results_folds_*文件夹中找到。6. 高级主题自定义数据集划分nnUNet 默认的 5 折交叉验证是随机划分的。但在某些场景下你可能需要固定训练集和验证集例如时间序列数据需要按时间划分。多中心数据需要按中心划分以避免数据泄露。与已有研究进行公平比较需要使用标准划分。这时你需要使用自定义数据集划分。6.1 创建自定义划分文件首先在nnUNet_preprocessed/Dataset120_MyLungCT/目录下创建一个名为splits_final.json的文件。其内容是一个列表列表中的每个元素是一个字典代表一种划分方案。[ { “train”: [“case_001”, “case_002”, “case_003”, “case_005”, “case_008”], “val”: [“case_004”, “case_006”, “case_007”] } ]train和val列表中的字符串是不带后缀和模态ID的案例名称。例如对应图像文件case_001_0000.png和标签文件case_001.png。这个例子只定义了一种划分一个字典。如果你想做 5 折交叉验证就需要在这个列表里放 5 个字典每个字典的train和val互斥且合起来是全部数据。6.2 使用自定义划分进行训练创建好splits_final.json后训练时指定 fold 为0即可nnUNetv2_train 120 2d 0nnUNet 会检测到splits_final.json文件并自动使用你定义的第一种划分方案列表中的第 0 个元素。如果你想训练第二种划分方案假设列表中有多个目前需要手动修改文件或使用其他 fold 索引但更常见的做法是只定义一种固定的划分。7. 评估与结果解读nnUNet 主要依赖 Dice 系数等指标在验证集上进行模型选择。对于测试集的正式评估你可以使用它提供的 Python API 或脚本。一个简单的评估方式是使用nnUNetv2_evaluate_simple命令行工具如果已安装或者直接运行其源码中的评估函数。评估的核心是计算预测结果与真实标签之间的重叠度指标。理解评估输出{ “foreground_mean”: { “Dice”: 0.905, “IoU”: 0.831, “TP”: 298.67, “FP”: 27.19, “FN”: 31.10 } }Dice: 最常用的医学图像分割指标范围 0-1越高越好。0.9 以上通常被认为是非常好的结果。IoU: 交并比与 Dice 高度相关但数值略低。TP/FP/FN: 真阳性、假阳性、假阴性的像素数量或区域数量可以帮助你分析模型是倾向于过分割FP高还是欠分割FN高。在整个流程跑通之后真正的迭代优化才刚刚开始。你可以回过头来审视数据质量检查标注的一致性可以尝试在dataset.json中调整patch_size的启发式参数甚至可以去修改nnUNetPlans.json中的网络深度、卷积核数量等超参数。nnUNetV2 提供的是一条高度自动化的基线而你的领域知识和对数据的理解才是让这条基线不断提升为 SOTA 解决方案的关键。我自己的经验是第一次成功跑通流程可能会遇到各种小问题但一旦打通后续对新数据集的适配速度会非常快这恰恰是 nnUNet 设计哲学的最大价值——将工程师从重复的工程劳动中解放出来。