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']:  # 分别处理两种模态