BraTS2025(nnUNet)

BraTS常用指令

预处理

数据处理要求文档
写个脚本先把MICCAI数据集文件的后缀转换成nnunet要求的格式,并且把文件夹的结构也换了,四个模态数据用模态映射处理
注意标签映射需要重新定义,有五个标签0-4,0是背景
后续在预处理的时候回报异常标签值
到时候用nibabel操作修改异常值为背景

之后要自定义数据集划分
先写个脚本把data_split改了,再模态映射

import os
import shutil
import re
import json

def create_dataset_json(target_dir, dataset_id, dataset_name):
    dataset_json = {
        "channel_names": {
            "0": "t1n",
            "1": "t1c",
            "2": "t2f",
            "3": "t2w"
        },
        "labels": {
            "background": 0,
            "NCR": 1,
            "NET": 2,
            "ED": 3,
            "ET": 4
        },
        "numTraining": 0,
        "file_ending": ".nii.gz"
    }
    
    json_path = os.path.join(target_dir, "dataset.json")
    with open(json_path, 'w') as f:
        json.dump(dataset_json, f, indent=4)
    print(f"dataset.json created at {json_path}")

def organize_brats_data(source_dir, target_dir, dataset_id, dataset_name):
    # 创建目标目录结构
    target_dataset_dir = os.path.join(target_dir, f"Dataset{dataset_id:03d}_{dataset_name}")
    images_tr_dir = os.path.join(target_dataset_dir, "imagesTr")
    labels_tr_dir = os.path.join(target_dataset_dir, "labelsTr")
    
    os.makedirs(images_tr_dir, exist_ok=True)
    os.makedirs(labels_tr_dir, exist_ok=True)
    
    # 创建 dataset.json
    create_dataset_json(target_dataset_dir, dataset_id, dataset_name)
    
    # 统计训练案例数量
    case_count = 0
    
    # 遍历源目录中的每个病例目录
    for case_dir in os.listdir(source_dir):
        case_dir_path = os.path.join(source_dir, case_dir)
        
        # 检查是否是目录
        if not os.path.isdir(case_dir_path):
            continue
        
        # 遍历病例目录中的文件
        for file_name in os.listdir(case_dir_path):
            file_path = os.path.join(case_dir_path, file_name)
            
            # 检查是否是文件
            if not os.path.isfile(file_path):
                continue
            
            # 匹配文件名格式 (BraTS-MET-XXXXX-XXX-*.nii.gz)
            match = re.match(r"(BraTS-MET-\d+-\d+)-([\w]+)\.nii\.gz", file_name)
            if not match:
                print(f"Unknown file format: {file_name}")
                continue
            
            case_identifier_part = match.group(1)
            modality = match.group(2)
            
            # 处理分割文件
            if modality == "seg":
                new_file_name = f"{case_identifier_part}.nii.gz"
                new_file_path = os.path.join(labels_tr_dir, new_file_name)
                shutil.move(file_path, new_file_path)
                print(f"Moved segmentation: {file_path} -> {new_file_path}")
                continue
            
            # 映射模态到四位数字标识符
            modality_id = None
            if modality == "t1n":
                modality_id = "0000"
            elif modality == "t1c":
                modality_id = "0001"
            elif modality == "t2f":
                modality_id = "0002"
            elif modality == "t2w":
                modality_id = "0003"
            
            if modality_id:
                new_file_name = f"{case_identifier_part}_{modality_id}.nii.gz"
                new_file_path = os.path.join(images_tr_dir, new_file_name)
                shutil.copy(file_path, new_file_path)
                print(f"Moved modality: {file_path} -> {new_file_path}")
            else:
                print(f"Unknown modality: {file_name}")
        
        case_count += 1
    
    # 更新 dataset.json 中的训练案例数量
    dataset_json_path = os.path.join(target_dataset_dir, "dataset.json")
    with open(dataset_json_path, 'r') as f:
        dataset_json = json.load(f)
    
    dataset_json["numTraining"] = case_count
    
    with open(dataset_json_path, 'w') as f:
        json.dump(dataset_json, f, indent=4)
    
    print(f"Successfully organized {case_count} cases into {target_dataset_dir}")

if __name__ == "__main__":
    source_dir = "/data/coding/nnUNet_raw"  # 源目录，包含 BraTS-MET-XXXXX-XXX 子目录
    target_dir = "/data/coding/nnUNet_raw"  # 目标目录，nnU-Net 的 raw 数据目录
    dataset_id = 1  # 数据集 ID
    dataset_name = "BraTS_MET"  # 数据集名称
    
    organize_brats_data(source_dir, target_dir, dataset_id, dataset_name)

生成一个split_final文件,之后训练的时候把折数设置为0 就可以自定义划分训练

预处理

先定义环境变量

export nnUNet_raw=/data/coding/nnUNet_raw
export nnUNet_preprocessed=/data/coding/nnUNet_preprocessed
export nnUNet_results=/data/coding/nnUNet_results

然后执行预处理命令

nnUNetv2_plan_and_preprocess -d 001 --verify_dataset_integrity
nnUNetv2_plan_and_preprocess -d 001 -c 2d --verify_dataset_integrity
nnUNetv2_plan_and_preprocess -d 001 -c 3d_fullres --verify_dataset_integrity

之后会自动检查数据集合法性,会有标签值异常的报错,用脚本修改异常值为背景

import os

import numpy as np

import nibabel as nib

  

def check_segmentation_labels(labels_folder: str, expected_labels: list) -> dict:

    """

    检查所有分割文件中的标签值，找出不符合预期的标签值。

  

    Args:

        labels_folder (str): 存储分割文件的目录路径。

        expected_labels (list): 预期的标签值列表。

  

    Returns:

        dict: 一个字典，键是文件名，值是不符合预期的标签值列表。

    """

    unexpected_labels = {}

    expected_labels_set = set(expected_labels)

  

    for label_file in os.listdir(labels_folder):

        label_path = os.path.join(labels_folder, label_file)

        label_data = nib.load(label_path).get_fdata()

        unique_labels = np.unique(label_data)

        unexpected = list(set(unique_labels) - expected_labels_set)

        if unexpected:

            unexpected_labels[label_file] = unexpected

  

    return unexpected_labels

  

# 使用示例

if __name__ == "__main__":

    labels_folder = "/data/coding/nnUNet_raw/Dataset001_BrainTumour/labelsTr"

    expected_labels = [0, 1, 2, 3, 4]

    unexpected = check_segmentation_labels(labels_folder, expected_labels)

    if unexpected:

        print("Found unexpected labels in the following files:")

        for file, labels in unexpected.items():

            print(f"{file}: {labels}")

    else:

        print("All segmentation files have expected labels.")

修改完之后就能预处理了

训练

2d命令

nnUNetv2_train Dataset001_BraTS_MET 2d 0

3d命令

nnUNetv2_train Dataset001_BrainTumour 3d_fullres 0

nohub命令

nohup sh -c 'export nnUNet_raw=/data/coding/nnUNet_raw && export nnUNet_preprocessed=/data/coding/nnUNet_preprocessed && export nnUNet_results=/data/coding/nnUNet_results && nnUNetv2_train Dataset001_BrainTumour 2d 0 --npz' > nnunet_train.log 2>&1 &

测试

nnUNetv2_predict -i /data/coding/nnUNet_raw/Dataset001_BraTS_MET/imagesTs -o /data/coding/nnUNet_results/predictions -d Dataset001_BraTS_MET -tr nnUNetTrainer -p nnUNetPlans -c 2d -f 0 --save_probabilities

可视化

nnunet输出的是.nii.gz文件,转成png比较费时间

import os
import json
import nibabel as nib
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from datetime import datetime

def create_consistent_colormap():
    """
    创建固定的颜色映射
    """
    colors = {
        0: [0, 0, 0],          # 背景-黑色
        1: [1, 0, 0],          # 标签1-红色
        2: [0, 1, 0],          # 标签2-绿色
        3: [0, 0, 1],          # 标签3-蓝色
        4: [1, 1, 0]           # 标签4-黄色
    }
    return colors

def visualize_slice(slice_data, output_path, color_dict):
    """
    可视化单个轴向切片，无边框和标记
    """
    # 创建图像但不设置大小，让图像填充整个空间
    plt.figure()
    
    # 创建RGB图像
    rgb_data = np.zeros((*slice_data.shape, 3))
    
    # 获取当前切片中存在的标签并着色
    for label in np.unique(slice_data):
        if label in color_dict:
            mask = (slice_data == label)
            for c in range(3):
                rgb_data[:, :, c][mask] = color_dict[label][c]
    
    # 显示图像，去除所有边框和空白
    plt.imshow(rgb_data)
    plt.axis('off')  # 关闭坐标轴
    
    # 移除所有边距
    plt.gca().set_position([0, 0, 1, 1])
    
    # 保存图像，确保无边框和边距
    plt.savefig(output_path, 
                bbox_inches='tight', 
                pad_inches=0,
                dpi=150)
    plt.close()

def load_test_samples(json_path):
    """
    从JSON文件中加载测试样本列表
    """
    with open(json_path, 'r') as f:
        data = json.load(f)
    return data.get('test', [])

def process_volume(file_path, output_dir, color_dict):
    """
    处理单个体积数据
    """
    # 加载NIfTI文件
    img = nib.load(str(file_path))
    data = img.get_fdata()
    data = np.round(data).astype(np.int32)
    
    # 处理所有轴向切片
    for i in range(data.shape[2]):
        slice_data = data[:, :, i]
        output_path = output_dir / f'axial_slice_{i:03d}.png'
        visualize_slice(slice_data, output_path, color_dict)

def process_predictions(test_samples, pred_dir, output_base_dir):
    """处理预测文件"""
    pred_dir = Path(pred_dir)
    output_base_dir = Path(output_base_dir)
    output_base_dir.mkdir(exist_ok=True, parents=True)
    
    color_dict = create_consistent_colormap()
    total_samples = len(test_samples)
    
    for idx, sample_name in enumerate(test_samples, 1):
        pred_file = pred_dir / f"{sample_name}.nii.gz"
        
        if not pred_file.exists():
            print(f"Warning: Prediction file not found: {pred_file}")
            continue
            
        try:
            print(f"Processing {idx}/{total_samples}: {sample_name}")
            output_dir = output_base_dir / sample_name
            output_dir.mkdir(exist_ok=True)
            
            process_volume(pred_file, output_dir, color_dict)
            
        except Exception as e:
            print(f"Error processing {pred_file}: {str(e)}")

def process_ground_truth(test_samples, gt_dir, output_base_dir):
    """处理真值文件"""
    gt_dir = Path(gt_dir)
    output_base_dir = Path(output_base_dir)
    output_base_dir.mkdir(exist_ok=True, parents=True)
    
    color_dict = create_consistent_colormap()
    total_samples = len(test_samples)
    
    for idx, sample_name in enumerate(test_samples, 1):
        gt_file = gt_dir / f"{sample_name}.nii.gz"
        
        if not gt_file.exists():
            print(f"Warning: Ground truth file not found: {gt_file}")
            continue
            
        try:
            print(f"Processing {idx}/{total_samples}: {sample_name}")
            output_dir = output_base_dir / sample_name
            output_dir.mkdir(exist_ok=True)
            
            process_volume(gt_file, output_dir, color_dict)
            
        except Exception as e:
            print(f"Error processing {gt_file}: {str(e)}")

def main():
    # 定义路径
    json_path = "/data/coding/test_data.json"
    pred_dir = "/data/coding/nnUNet_results/predictions"
    gt_dir = "/data/coding/nnUNet_raw/Dataset001_BraTS_MET/labelsTr"
    pred_output_dir = "/data/coding/vis"
    gt_output_dir = "/data/coding/vis_GT"
    
    # 加载测试样本列表
    test_samples = load_test_samples(json_path)
    print(f"Found {len(test_samples)} test samples")
    
    # 处理预测结果和真值
    print("\nProcessing predictions...")
    process_predictions(test_samples, pred_dir, pred_output_dir)
    
    print("\nProcessing ground truth...")
    process_ground_truth(test_samples, gt_dir, gt_output_dir)

if __name__ == "__main__":
    main()