nnUNet实战：用MMWHS心脏数据集完成3D医学影像分割（附完整代码）

nnUNet实战：用MMWHS心脏数据集完成3D医学影像分割

医学影像分割一直是计算机辅助诊断中的核心任务，而心脏结构的精确分割对于心血管疾病诊断和治疗规划尤为重要。MMWHS（Multi-Modality Whole Heart Segmentation）数据集作为心脏影像分割领域的标杆数据集，包含了CT和MR两种模态的完整心脏扫描数据。本文将带你从零开始，使用nnUNet框架完成MMWHS数据集的3D心脏结构分割任务。

1. 环境配置与数据准备

在开始之前，我们需要搭建适合nnUNet运行的环境。nnUNet对PyTorch版本有特定要求，建议使用PyTorch 1.9+和CUDA 11.1+的组合以获得最佳性能。

# 创建并激活conda环境
conda create -n nnunet python=3.8
conda activate nnunet

# 安装PyTorch（根据CUDA版本选择）
pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116

# 安装nnUNet v2
pip install nnunetv2

接下来设置nnUNet必需的三个环境变量，这些路径需要有足够的存储空间（建议至少500GB）：

# 编辑~/.bashrc文件，添加以下内容
export nnUNet_raw="/path/to/nnUNet_raw"  # 原始数据集存储路径
export nnUNet_preprocessed="/path/to/nnUNet_preprocessed"  # 预处理数据路径
export nnUNet_results="/path/to/nnUNet_results"  # 训练结果保存路径

# 使环境变量生效
source ~/.bashrc

从MMWHS官网下载数据集后，你会得到四个文件夹：

• mr_train：MR训练图像
• mr_test：MR测试图像
• ct_train：CT训练图像
• ct_test：CT测试图像

2. 数据集转换与多模态处理

MMWHS数据集需要转换为nnUNet要求的格式。由于包含CT和MR两种模态，我们需要特别注意数据对齐问题。以下是完整的转换脚本：

import os
import shutil
from pathlib import Path
from batchgenerators.utilities.file_and_folder_operations import *
from nnunetv2.dataset_conversion.generate_dataset_json import generate_dataset_json

def convert_mmwhs(src_data_folder: str, dataset_id: int = 27):
    # 创建输出目录结构
    out_dir = Path(os.environ['nnUNet_raw']) / f"Dataset{dataset_id:03d}_MMWHS"
    (out_dir / "imagesTr").mkdir(parents=True, exist_ok=True)
    (out_dir / "labelsTr").mkdir(parents=True, exist_ok=True)
    (out_dir / "imagesTs").mkdir(parents=True, exist_ok=True)

    # 处理MR模态数据
    num_training_cases = 0
    for modality in ['mr', 'ct']:  # 分别处理两种模态
        train_dir = Path(src_data_folder) / f"{modality}_train"
        test_dir = Path(src_data_folder) / f"{modality}_test"
        
        # 处理训练集
        for file in train_dir.glob("*.nii.gz"):
            if "_label" in file.name:  # 标签文件
                shutil.copy(file, out_dir / "labelsTr" / file.name.replace("_label", ""))
            else:  # 图像文件
                modality_id = "0000" if modality == 'mr' else "0001"
                shutil.copy(file, out_dir / "imagesTr" / f"{file.stem.split('.')[0]}_{modality_id}.nii.gz")
                num_training_cases += 1
        
        # 处理测试集
        for file in test_dir.glob("*.nii.gz"):
            if "_label" not in file.name:  # 只复制图像文件
                modality_id = "0000" if modality == 'mr' else "0001"
                shutil.copy(file, out_dir / "imagesTs" / f"{file.stem.split('.')[0]}_{modality_id}.nii.gz")

    # 生成dataset.json文件
    generate_dataset_json(
        str(out_dir),
        channel_names={
            "0": "MR",
            "1": "CT"
        },
        labels={
            "background": 0,
            "LV myocardium": 1,
            "LV cavity": 2,
            "RV cavity": 3,
            "left atrium": 4,
            "right atrium": 5,
            "ascending aorta": 6,
            "pulmonary artery": 7
        },
        file_ending=".nii.gz",
        num_training_cases=num_training_cases,
        overwrite_image_reader_writer='NibabelIOWithReorient'
    )

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input_folder", required=True, help="MMWHS数据集目录")
    parser.add_argument("-d", "--dataset_id", type=int, default=27, help="nnUNet数据集ID")
    args = parser.parse_args()
    convert_mmwhs(args.input_folder, args.dataset_id)

注意：MMWHS数据集中不同模态的图像分辨率不一致，直接融合可能导致问题。我们的解决方案是为每种模态创建独立的数据集，后续训练时使用nnUNet的交叉模态训练策略。

执行转换脚本后，检查生成的数据集结构：

Dataset027_MMWHS_MR/
├── imagesTr/
│   ├── case1_0000.nii.gz  # MR图像
│   └── ...
├── labelsTr/
│   ├── case1.nii.gz
│   └── ...
└── dataset.json

Dataset028_MMWHS_CT/
├── imagesTr/
│   ├── case1_0000.nii.gz  # CT图像 
│   └── ...
├── labelsTr/
│   ├── case1.nii.gz
│   └── ...
└── dataset.json

3. 数据预处理与实验规划

nnUNet会自动执行一系列复杂的预处理步骤，包括：

• 重采样到统一分辨率
• 标准化（CT使用Hounsfield单位截断，MR使用z-score标准化）
• 数据增强（旋转、缩放、弹性变形等）

运行以下命令启动预处理：

# 分别处理MR和CT数据集
nnUNetv2_plan_and_preprocess -d 27 --verify_dataset_integrity
nnUNetv2_plan_and_preprocess -d 28 --verify_dataset_integrity

预处理完成后，可以在nnUNet_preprocessed目录下看到生成的计划文件，例如：

# plans_3D.json 关键内容
{
    "data_identifier": "nnUNetPlans_3D",
    "spacing": [1.5, 1.5, 1.5],  # 统一分辨率
    "stages": [
        {
            "patch_size": [128, 128, 128],  # 3D全分辨率使用的patch大小
            "batch_size": 2,
            "num_pool_per_axis": [5, 5, 5],
            "do_dummy_2D_data_aug": false
        }
    ],
    "normalization_schemes": ["ZScore"],  # MR使用Z-score
    "ct_properties": {  # CT特有的设置
        "lower_bound": -950,
        "upper_bound": 1750
    }
}

4. 模型训练与调优

nnUNet默认会训练五种配置（2D、3D全分辨率、3D低分辨率、级联等），我们重点使用3D全分辨率配置：

# 训练MR模型
nnUNetv2_train 27 3d_fullres 0 -tr nnUNetTrainer --npz

# 训练CT模型（使用不同的学习率）
nnUNetv2_train 28 3d_fullres 0 -tr nnUNetTrainer --npz --lr 0.0003

提示：使用CUDA_VISIBLE_DEVICES参数指定GPU，如CUDA_VISIBLE_DEVICES=0,1使用前两块GPU。

训练过程中可能遇到的常见问题及解决方案：

问题现象	可能原因	解决方案
OOM错误	batch size过大	修改plans.json减小batch size或patch size
损失不下降	学习率不合适	添加--lr参数调整学习率
Dice波动大	数据分布不均	检查标签是否连续，背景类是否为0

对于多模态数据，我们可以使用nnUNet的交叉模态训练策略：

# 修改nnUNetTrainer.py中的train_step方法
def train_step(self, batch):
    # 对MR和CT数据分别应用不同的标准化
    if self.is_mr:
        data = (batch['data'] - self.mr_mean) / self.mr_std
    else:
        data = torch.clamp(batch['data'], self.ct_lower, self.ct_upper)
        data = (data - self.ct_mean) / self.ct_std
    
    # 后续网络前向计算...
    return super().train_step({'data': data, 'target': batch['target']})

5. 模型推理与结果可视化

训练完成后，使用以下命令进行预测：

# 查找最佳配置
nnUNetv2_find_best_configuration 27 -c 3d_fullres

# 执行推理
nnUNetv2_predict -i /input_path -o /output_path -d 27 -c 3d_fullres -f all

结果可视化可以使用以下Python代码：

import nibabel as nib
import matplotlib.pyplot as plt
import numpy as np

def visualize_results(image_path, pred_path, label_path=None, slice_idx=None):
    # 加载数据
    image = nib.load(image_path).get_fdata()
    pred = nib.load(pred_path).get_fdata()
    
    if label_path:
        label = nib.load(label_path).get_fdata()
    
    # 自动选择最佳切片
    if slice_idx is None:
        slice_idx = image.shape[2] // 2
    
    # 创建可视化
    fig, axes = plt.subplots(1, 3 if label_path else 2, figsize=(15, 5))
    axes[0].imshow(image[:, :, slice_idx], cmap='gray')
    axes[0].set_title('Input Image')
    
    axes[1].imshow(pred[:, :, slice_idx], cmap='jet', alpha=0.5)
    axes[1].set_title('Prediction')
    
    if label_path:
        axes[2].imshow(label[:, :, slice_idx], cmap='jet', alpha=0.5)
        axes[2].set_title('Ground Truth')
    
    plt.tight_layout()
    plt.show()

# 示例使用
visualize_results(
    "imagesTs/case1_0000.nii.gz",
    "predictions/case1.nii.gz",
    "labelsTs/case1.nii.gz"
)

对于定量评估，nnUNet提供了内置的评估工具：

nnUNetv2_evaluate_folder -ref /path/to/ground_truth -pred /path/to/predictions -l 1 2 3 4 5 6 7

典型的心脏结构分割性能指标：

结构	Dice系数	HD95(mm)	体积相关性
LV心肌	0.92±0.03	2.1±0.8	0.98
LV腔	0.94±0.02	1.8±0.6	0.99
升主动脉	0.88±0.05	2.5±1.2	0.97

6. 高级技巧与性能优化

多GPU训练加速：

# 使用DDP进行多GPU训练
torchrun --nnodes=1 --nproc_per_node=4 --master_port=1234 \
    nnUNetv2_run_training DDP 27 3d_fullres 0 -tr nnUNetTrainer --npz

混合精度训练： 在nnUNetTrainer中添加：

from torch.cuda.amp import autocast, GradScaler

class nnUNetTrainerAMP(nnUNetTrainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.scaler = GradScaler()
    
    def train_step(self, batch):
        with autocast():
            loss = super().train_step(batch)
        self.scaler.scale(loss).backward()
        self.scaler.step(self.optimizer)
        self.scaler.update()
        return loss

模型集成：

# 训练多个fold并集成
for fold in 0 1 2 3 4:
    nnUNetv2_train 27 3d_fullres $fold
    
nnUNetv2_ensemble -f /path/to/fold0 /path/to/fold1 /path/to/fold2 -o /ensemble_output

实际项目中，在NVIDIA A100上训练完整MMWHS数据集（20例）的典型耗时：

配置	耗时	GPU内存占用
3D全分辨率	18小时	40GB
2D配置	6小时	12GB
级联模型	24小时	24GB+40GB

通过本文的实践指南，你应该已经掌握了使用nnUNet处理多模态心脏影像分割的全流程。这套方法同样适用于其他器官的3D医学影像分割任务，只需按照相同的数据准备流程即可快速适配。