一文讲清nnUNetv2在自己数据集上训练

本文只讲述nnunetv2的在2D图像的上复现步骤，对于实现细节可以阅读原文和代码！

1. 下载代码与构建环境：

git clone https://github.com/MIC-DKFZ/nnUNet.git  # 下载代码
cd nnUNet  # 切换目录
conda create -n myenv python=3.9  # 注意nnUNetv2需要python>=3.9
conda activate myenv
pip install nnunetv2
pip install -e .  #最后这个点也不能忽略

安装隐藏层（可选，可以不安装），hiddenlayer 使 nnU-net 能够生成网络拓扑图

pip install --upgrade git+https://github.com/FabianIsensee/hiddenlayer.git

2. 数据准备

（1）在nnunet 文件夹下创建DATASET文件夹，并创建 nnUNet_raw、nnUNet_preprocessed 和 nnUNet_trained_models，文件目录如下：

3.设置环境路径

（1）在终端（用户根目录路径）输入：

echo 'export nnUNet_raw="path/nnUNet/DATASET/nnUNet_raw"' >> ~/.bashrc
echo 'export nnUNet_preprocessed="path/nnUNet/DATASET/nnUNet_preprocessed"' >> ~/.bashrc
echo 'export nnUNet_results="path/nnUNet/DATASET/nnUNet_trained_models"' >> ~/.bashrc

注意修改“path”为自己的路径

（2）终端运行代码更新 .bashrc:

source ~/.bashrc

================================= 或者 =================================
可修改文件：“…/nnUNet/nnunetv2/paths.py”

nnUNet_raw = os.environ.get('根目录/nnUNet/DATASET/nnUNet_raw')
nnUNet_preprocessed = os.environ.get('根目录/nnUNet/DATASET/nnUNet_preprocessed')
nnUNet_results = os.environ.get('根目录/nnUNet/DATASET/nnUNet_trained_models')

如遇到如下问题：
nnUNet_raw is not defined and nnU-Net can only be used on data for which preprocessed files " "are already present on your system. nnU-Net cannot be used for experiment planning and preprocessing like " "this. If this is not intended, please read documentation/setting_up_paths.md for information on how to set " "this up properly.
…
则修改代码如下：

nnUNet_raw="根目录/nnUNet/DATASET/nnUNet_raw"
nnUNet_preprocessed="根目录/nnUNet/DATASET/nnUNet_preprocessed"
nnUNet_results="根目录/nnUNet/DATASET/nnUNet_trained_models"

4. 转换数据集

（1）将自己的数据集创建为如下结构：

input: 图像数据
output:标签数据
（2）修改文件“…/nnUNet/nnunetv2/dataset_conversion/Dataset120_RoadSegmentation.py”：

import multiprocessing
import shutil
from multiprocessing import Pool

from batchgenerators.utilities.file_and_folder_operations import *

from nnunetv2.dataset_conversion.generate_dataset_json import generate_dataset_json
from nnunetv2.paths import nnUNet_raw
from skimage import io
from acvl_utils.morphology.morphology_helper import generic_filter_components
from scipy.ndimage import binary_fill_holes


def load_and_covnert_case(input_image: str, input_seg: str, output_image: str, output_seg: str,
                          min_component_size: int = 50):
    seg = io.imread(input_seg)
    # seg[seg == 255] = 1
    seg[seg == 18] = 1
    seg[seg == 54] = 2
    seg[seg == 255] = 0
    image = io.imread(input_image)
    image = image.sum(2)
    mask = image == (3 * 255)
    # the dataset has large white areas in which road segmentations can exist but no image information is available.
    # Remove the road label in these areas
    mask = generic_filter_components(mask, filter_fn=lambda ids, sizes: [i for j, i in enumerate(ids) if
                                                                         sizes[j] > min_component_size])
    mask = binary_fill_holes(mask)
    seg[mask] = 0
    io.imsave(output_seg, seg, check_contrast=False)
    shutil.copy(input_image, output_image)


if __name__ == "__main__":
    # extracted archive from https://www.kaggle.com/datasets/insaff/massachusetts-roads-dataset?resource=download
    # source = '/media/fabian/data/raw_datasets/Massachussetts_road_seg/road_segmentation_ideal'
    source = '../nnunet_data'
    nnUNet_raw = '../nnUNet/DATASET/nnUNet_raw/'

    # dataset_name = 'Dataset120_RoadSegmentation'
    dataset_name = 'Dataset001_XXX'

    imagestr = join(nnUNet_raw, dataset_name, 'imagesTr')
    imagests = join(nnUNet_raw, dataset_name, 'imagesTs')
    labelstr = join(nnUNet_raw, dataset_name, 'labelsTr')
    labelsts = join(nnUNet_raw, dataset_name, 'labelsTs')
    maybe_mkdir_p(imagestr)
    maybe_mkdir_p(imagests)
    maybe_mkdir_p(labelstr)
    maybe_mkdir_p(labelsts)

    train_source = join(source, 'training')
    test_source = join(source, 'testing')

    with multiprocessing.get_context("spawn").Pool(8) as p:

        # not all training images have a segmentation
        valid_ids = subfiles(join(train_source, 'output'), join=False, suffix='png')
        num_train = len(valid_ids)
        r = []
        for v in valid_ids:
            r.append(
                p.starmap_async(
                    load_and_covnert_case,
                    ((
                         join(train_source, 'input', v),
                         join(train_source, 'output', v),
                         join(imagestr, v[:-4] + '_0000.png'),
                         join(labelstr, v),
                         50
                     ),)
                )
            )

        # test set
        valid_ids = subfiles(join(test_source, 'output'), join=False, suffix='png')
        for v in valid_ids:
            r.append(
                p.starmap_async(
                    load_and_covnert_case,
                    ((
                         join(test_source, 'input', v),
                         join(test_source, 'output', v),
                         join(imagests, v[:-4] + '_0000.png'),
                         join(labelsts, v),
                         50
                     ),)
                )
            )
        _ = [i.get() for i in r]

    # generate_dataset_json(join(nnUNet_raw, dataset_name), {0: 'R', 1: 'G', 2: 'B'}, {'background': 0, 'road': 1},
    #                       num_train, '.png', dataset_name=dataset_name)
    generate_dataset_json(join(nnUNet_raw, dataset_name), {0: 'R', 1: 'G', 2: 'B'}, {'background': 0, 'class1': 1, 'class2': 2},
                      num_train, '.png', dataset_name=dataset_name)

以下部分根据自己的任务修改：

seg[seg == 18] = 1
seg[seg == 54] = 2
seg[seg == 255] = 0

source = '../nnunet_data'
nnUNet_raw = '../nnUNet/DATASET/nnUNet_raw/'
dataset_name = 'Dataset001_XXX'

（3）运行Dataset120_RoadSegmentation.py文件，即可在"…/nnUNet/DATASET/nnUNet_raw/Dataset001_XXX"下生成对应的训练格式和dataset.json文件,如下图所示：

5. 数据集预处理

此步骤对数据进行：裁剪crop，重采样resample以及标准化normalization

nnUNetv2_plan_and_preprocess -d DATASET_ID --verify_dataset_integrity  

注： DATASET_ID是001 (Dataset001_XXX)
会在 nnUNet_preprocessed 文件夹中创建一个以数据集命名的新子文件夹。命令完成后，将出现一个 dataset_fingerprint.json 文件以及一个 nnUNetPlans.json 文件。还有一些子文件夹包含 UNet 配置的预处理数据。

5.训练命令

CUDA_VISIBLE_DEVICES=1 nnUNetv2_train 1 2d 0
CUDA_VISIBLE_DEVICES=1 nnUNetv2_train 1 2d 1
CUDA_VISIBLE_DEVICES=1 nnUNetv2_train 1 2d 2
CUDA_VISIBLE_DEVICES=1 nnUNetv2_train 1 2d 3
CUDA_VISIBLE_DEVICES=1 nnUNetv2_train 1 2d 4

‘CUDA_VISIBLE_DEVICES=1’ 表示指定GPU训练
‘2d’ 是选用2D Unet模型
‘1’ 表示任务编码，Dataset001_XXX
‘0,1,2,3,4’ 代表五折交叉验证

若运行中断，代码会50个epoch保存一下checkpoint ，若中断使用：原来指令后面加–c可以接着运行，用运行第0折时中断为例：

CUDA_VISIBLE_DEVICES=1 nnUNetv2_train 1 2d 0 --c

6. 推理结果

nnUNetv2_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -d DATASET_NAME_OR_ID -c CONFIGURATION --save_probabilities
# nnUNetv2_predict -h 查看参数解析

INPUT_FOLDER：测试数据地址
OUTPUT_FOLDER：分割结果存放地址
DATASET_NAME_OR_ID：数据集ID
CONFIGURATION：使用的什么架构，2d or 3d_fullres or 3d_cascade_fullres等，这里训练用的什么就写什么
save_probabilities：将预测概率与需要大量磁盘空间的预测分段掩码一起保存。
默认情况下，推理将通过交叉验证的所有 5 个折叠作为一个整体来完成（根据5个模型得到一个结果）。我们强烈建议您使用全部 5 折。因此，在运行推理之前必须训练所有 5 个折叠。
要想每个模型分开得到结果，就加参数-f all，或者只为某一折出结果，就加参数-f 2(得到 fold 2 的结果)

以fold 2 举例：

nnUNetv2_predict -i input -o output -d 001 -c 3d_lowres -f 2

注意： 如需要在其他的数据集测试需要进行图像预处理工作，具体参考步骤 4, 5.

7. 获得最佳配置

nnUNetv2_find_best_configuration DATASET_NAME_OR_ID -c CONFIGURATIONS 

DATASET_NAME_OR_ID： 数据集ID，即001—Dataset001
CONFIGURATIONS：2d
以2d为例，任务为001，折数为1（不加折数，默认从五折中找最优的），则命令为：

nnUNetv2_find_best_configuration 001 -c 2d  -f 1

则会在路径“~/nnUNet/DATASET/nnUNet_trained_models/Dataset001_XXX" 目录下生成：inference_information.json和inference_instructions.txt文件：
inference_information.json：保存的最佳的推理结果
inference_instructions.txt：保存的推理命令

nnUNetv2_predict -d Dataset001_XXX -i INPUT_FOLDER -o OUTPUT_FOLDER -f  0 1 2 3 4 -tr nnUNetTrainer -c 2d -p nnUNetPlans

8. 对无标签数据进行推理

** 即已经有该数据集训练好的模型，需要对新的图片（无标签）进行推理。**

（1）将图片命名为nnunet需要的格式xxx_0000.png（示例代码如下）：

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import os
from pathlib import Path

def rename_images(input_dir):
    """
    Rename image files from X.png to X_0000.png format
    
    Args:
        input_dir: Directory containing the image files
    """
    input_path = Path(input_dir)
    
    # Check if directory exists
    if not input_path.exists():
        raise RuntimeError(f"Input directory {input_dir} does not exist")
        
    # Get all PNG files
    png_files = list(input_path.glob("*.png"))
    
    for file_path in png_files:
        # Get original filename without extension
        filename = file_path.stem
        
        # Create new filename
        new_filename = f"{filename}_0000.png"
        new_path = file_path.parent / new_filename
        
        # Rename file
        file_path.rename(new_path)
        print(f"Renamed {file_path.name} to {new_filename}")

if __name__ == "__main__":
    # Specify input directory
    input_dir = "your_file_path"
    
    # Process the renaming
    rename_images(input_dir)
    print("Renaming completed!")

（2）使用 “ 6. 推理结果 ” 的设置进行推理即可。

参考资料：

1. https://blog.csdn.net/weixin_45882172/article/details/138216705?spm=1001.2101.3001.6650.5&utm_medium=distribute.pc_relevant.none-task-blog-2%7Edefault%7Ebaidujs_baidulandingword%7ECtr-5-138216705-blog-131539779.235%5Ev43%5Epc_blog_bottom_relevance_base6&depth_1-utm_source=distribute.pc_relevant.none-task-blog-2%7Edefault%7Ebaidujs_baidulandingword%7ECtr-5-138216705-blog-131539779.235%5Ev43%5Epc_blog_bottom_relevance_base6&utm_relevant_index=10