nnUNetv2推理部分||源码解读||深度学习||医学图像处理

前言：nnUNet推理部分解读，主要包括整体设计理念、推理部分各个函数以及相关参数、推理流程、一些函数的局部细节以及涉及到的一些知识点和思想。

一.设计理念

先理解下面这两个设计理念，对于阅读源码会有很大的帮助。

1.面向过程和面向对象编程思想

对于习惯使用Python脚本解决问题的开发者来说，通常采用面向过程的思想。这种思想强调通过自顶向下的方式设计程序，将复杂问题分解为一系列步骤。打个比方，如果要自己实现一个nnUNet推理部分，基于面向过程的思想通常会考虑以下步骤：
1.读取数据：根据输入路径，加载需要预测的数据集。
2.数据预处理：对数据进行转换，使其满足模型输入的格式要求。
3.模型推理：将处理后的数据输入模型，获取预测结果。
可以看到的是，这种思路关注的是过程和函数，强调解决问题的步骤，通过自顶向下的方式来解决问题。
但是，对于大型项目而言，这样的设计思路可能存在一些问题。例如，数据集的预处理可能在模型训练前就已经完成，没有必要在推理阶段重复实现。因此，需要结合面向对象的思想，将预处理、训练和推理分别抽象为不同的对象。面向对象的思想关注的是对象，将问题域中的实体抽象为对象，强调对象之间的交互。程序被视为一组相互协作的对象，通过类和对象来组织代码，自底向上地解决问题。下面是我认为的一些关键点：
1.架构设计：首先，从整体架构的角度审视代码。考虑代码的组织结构、模块划分以及它们之间的交互方式。
2.功能定义：其次，明确需要实现的功能，理清各个功能之间的关系和数据流向。考虑系统的输入、输出以及中间处理过程，并确定所需的参数。
3.代码质量：最后，关注代码的鲁棒性和简洁性。进行必要的错误检查和异常处理，同时将重复的功能封装成模块或函数，提高代码的复用性。
面向过程和面向对象是两种不同的编程思想，各有各的用武之地。简单的脚本任务，用面向过程就很直观高效；但如果是复杂的大项目，面向对象就能带来更好的代码组织、可维护性和可扩展性。

2.nnUNet整体的设计理念

相较于其他深度学习模型框架，nnU-Net 的主要区别在于其高度自动化和数据驱动的设计理念。nnU-Net 会通过分析医学图像的空间信息（如尺寸、体素间距）、模态（如 CT、MRI）、标签（如类别数）等特征，自动生成相应的预处理步骤、网络结构以及后处理方案。
自动化生成网络架构与超参数配置
在传统的深度学习框架中，网络的架构和超参数通常是通过配置文件（如 YAML 或 JSON）手动设置，或直接硬编码到程序中。这种方式尽管直观，但灵活性有限。当更换数据集或任务时，往往需要手动修改配置文件或代码，增加了开发复杂度。
nnUNet 则采用了一种更为智能的自动化方法。它通过分析数据集的内在属性（如图像尺寸、通道数、类别数等），自动生成最适合该数据集的网络架构和超参数。这一过程被称为数据集指纹（dataset fingerprint）。在此过程中，nnUNet 会根据数据集的特点提取关键信息，并存储在 JSON 文件中，如 dataset.json 和 plans.json。
动态构建与加载
在模型的训练和推理阶段，nnUNet 通过读取这些 JSON 文件，动态地构建网络结构并加载相关配置。这种基于数据的自动化设计使得 nnUNet 能够自适应不同的数据集，避免了手动调整架构和超参数的繁琐过程。相较于其他框架，nnUNet 不需要用户具备深厚的模型设计经验，能够根据不同的任务自动优化配置，极大地简化了使用门槛。但正因如此，如果想要自定义网络结构，则比其他框架更为复杂。
开箱即用与错误预防
另一个显著的优势是其开箱即用的设计理念。用户只需提供数据，nnUNet 就能够全程自动化地处理数据预处理、模型设计、训练和测试，用户无需深入了解或手动设置大量参数。正因如此，nnUNet 在参数设置上进行了严格的检查，尤其针对输入数据的边界条件，减少了由于配置不当而产生的错误。在阅读源码时可以看到大量的检查和判断。

二.推理部分主要函数以及相关参数

1. 导入模块

代码首先导入了许多必要的 Python 模块和 nnUNet 框架中的自定义模块，这些模块提供了数据处理、模型加载、推理、多线程处理等功能。
根据前面提到的整体设计理念，大概就可以知道为什么nnUNet推理部分会导入大量的自定义模块了。

2. nnUNetPredictor 类

nnUNetPredictor 类是用于执行推理的核心类，它包含了初始化模型、加载模型权重、执行推理等方法。

2.1 __init__ 方法

• 参数:
  • tile_step_size: 滑动窗口预测的步长。
  • use_gaussian: 是否在滑动窗口预测中使用高斯加权。
  • use_mirroring: 是否在推理过程中使用镜像数据增强。
  • perform_everything_on_device: 是否在设备上执行所有操作（如 GPU）。
  • device: 指定使用的设备（如 cuda 或 cpu）。
  • verbose: 是否打印详细信息。
  • verbose_preprocessing: 是否打印预处理详细信息。
  • allow_tqdm: 是否使用 tqdm 显示进度条。
• 初始化:
  • 初始化了一些实例变量，如 self.plans_manager, self.configuration_manager, self.network 等。
  • 如果设备类型是 cuda，则启用 cudnn 的基准模式以加速计算。

2.2 initialize_from_trained_model_folder 方法

• 功能: 从训练好的模型文件夹中初始化推理器。
• 参数:
  • model_training_output_dir: 训练输出目录。
  • use_folds: 使用的折数。
  • checkpoint_name: 检查点文件名。
• 步骤:
  1. 自动检测可用的折数。
  2. 加载 dataset.json 和 plans.json 文件。
  3. 根据检查点文件加载模型权重和配置信息。
  4. 构建网络架构并加载权重。
  5. 初始化 nnUNetPredictor 的实例变量。

2.3 manual_initialization 方法

• 功能: 手动初始化推理器，通常用于训练过程中的验证。
• 参数:
  • network: 神经网络模型。
  • plans_manager: 计划管理器。
  • configuration_manager: 配置管理器。
  • parameters: 模型参数。
  • dataset_json: 数据集 JSON 文件。
  • trainer_name: 训练器名称。
  • inference_allowed_mirroring_axes: 允许的镜像轴。

2.4 auto_detect_available_folds 方法

• 功能: 自动检测可用的折数。
• 参数:
  • model_training_output_dir: 训练输出目录。
  • checkpoint_name: 检查点文件名。

2.5 _manage_input_and_output_lists 方法

• 功能: 管理输入和输出文件列表。
• 参数:
  • list_of_lists_or_source_folder: 输入文件列表或源文件夹。
  • output_folder_or_list_of_truncated_output_files: 输出文件夹或输出文件列表。
  • folder_with_segs_from_prev_stage: 前一阶段的分割结果文件夹。
  • overwrite: 是否覆盖已存在的文件。
  • part_id: 当前部分的 ID。
  • num_parts: 总部分数。
  • save_probabilities: 是否保存概率。

2.6 predict_from_files 方法

• 功能: 从文件中进行预测。
• 参数:
  • list_of_lists_or_source_folder: 输入文件列表或源文件夹。
  • output_folder_or_list_of_truncated_output_files: 输出文件夹或输出文件列表。
  • save_probabilities: 是否保存概率。
  • overwrite: 是否覆盖已存在的文件。
  • num_processes_preprocessing: 预处理进程数。
  • num_processes_segmentation_export: 分割结果导出进程数。
  • folder_with_segs_from_prev_stage: 前一阶段的分割结果文件夹。
  • num_parts: 总部分数。
  • part_id: 当前部分的 ID。

2.7 _internal_get_data_iterator_from_lists_of_filenames 方法

• 功能: 从文件名列表中获取数据迭代器。
• 参数:
  • input_list_of_lists: 输入文件列表。
  • seg_from_prev_stage_files: 前一阶段的分割结果文件列表。
  • output_filenames_truncated: 输出文件名列表。
  • num_processes: 进程数。

2.8 get_data_iterator_from_raw_npy_data 方法

• 功能: 从原始的 Numpy 数组中获取数据迭代器。
• 参数:
  • image_or_list_of_images: 图像或图像列表。
  • segs_from_prev_stage_or_list_of_segs_from_prev_stage: 前一阶段的分割结果或结果列表。
  • properties_or_list_of_properties: 属性或属性列表。
  • truncated_ofname: 输出文件名或文件名列表。
  • num_processes: 进程数。

2.9 predict_from_list_of_npy_arrays 方法

• 功能: 从 Numpy 数组列表中进行预测。
• 参数:
  • image_or_list_of_images: 图像或图像列表。
  • segs_from_prev_stage_or_list_of_segs_from_prev_stage: 前一阶段的分割结果或结果列表。
  • properties_or_list_of_properties: 属性或属性列表。
  • truncated_ofname: 输出文件名或文件名列表。
  • num_processes: 进程数。
  • save_probabilities: 是否保存概率。
  • num_processes_segmentation_export: 分割结果导出进程数。

2.10 predict_from_data_iterator 方法

• 功能: 从数据迭代器中进行预测。
• 参数:
  • data_iterator: 数据迭代器。
  • save_probabilities: 是否保存概率。
  • num_processes_segmentation_export: 分割结果导出进程数。

2.11 predict_single_npy_array 方法

• 功能: 对单个 Numpy 数组进行预测。
• 参数:
  • input_image: 输入图像。
  • image_properties: 图像属性。
  • segmentation_previous_stage: 前一阶段的分割结果。
  • output_file_truncated: 输出文件名。
  • save_or_return_probabilities: 是否保存或返回概率。

2.12 predict_logits_from_preprocessed_data 方法

• 功能: 从预处理数据中预测 logits。
• 参数:
  • data: 预处理数据。

2.13 _internal_get_sliding_window_slicers 方法

• 功能: 获取滑动窗口切片器。
• 参数:
  • image_size: 图像尺寸。

2.14 _internal_maybe_mirror_and_predict 方法

• 功能: 可能进行镜像并预测。

• 参数:

  • x: 输入数据。

2.15 _internal_predict_sliding_window_return_logits 方法

• 功能: 内部滑动窗口预测并返回 logits。
• 参数:
  • data: 输入数据。
  • slicers: 切片器。
  • do_on_device: 是否在设备上执行。

2.16 predict_sliding_window_return_logits 方法

• 功能: 滑动窗口预测并返回 logits。
• 参数:
  • input_image: 输入图像。

3. 入口点函数

代码还提供了两个入口点函数 predict_entry_point_modelfolder 和 predict_entry_point，用于从命令行运行推理。

4. 主函数

在 __main__ 中，代码展示了如何使用 nnUNetPredictor 类进行推理，包括从文件夹中预测和从 Numpy 数组中预测。

三.推理流程

1.官方文档

这里可以先看官方给出的README.md，这里给出翻译版本：

nnU-Net 的推理现在比之前更加动态，使得你可以更无缝地将 nnU-Net 集成到现有的工作流程中。 本 README 将为你快速概述这些选项。这并不是一份完整的指南，要了解所有细节请查看代码！
前言
在速度方面，最有效的推理策略是 nnU-Net 默认的推理方式！图像会被实时读取，并在后台工作线程中进行预处理。主进程会获取预处理后的图像，进行预测，并将预测结果发送给另一组后台工作线程，后者会调整输出的 logits 大小，将它们转换为分割结果并导出分割。
默认设置是最佳选择的原因如下：
1.图像加载和预处理以及分割结果导出与预测是交错进行的。主进程可以专注于与计算设备（即 GPU）的通信，而不需要执行其他处理。这最大化地利用了你的资源！
2.仅存储当前需要的图像和分割结果在内存中！想象一下，如果需要预测很多图像并且还要将所有图像和结果都存储在系统内存中，将会极大地占用资源。
nnUNetPredictor
新的 nnUNetPredictor 类封装了推理代码，使得在不同模式之间切换变得更加简单。你的代码可以持有一个 nnUNetPredictor 实例并实时进行预测。之前这是无法实现的，每次新的预测请求都会重新加载参数并重新实例化网络架构，这并不理想。
nnUNetPredictor 必须手动初始化！在 99% 的使用场景中，你需要使用 predictor.initialize_from_trained_model_folder 函数。
注意：如果你没有指定输出文件夹或输出文件，则预测的分割结果将会直接返回。

A.推荐的 nnU-Net 默认设置：从源文件进行预测
简要说明：

实时加载图像
在后台工作线程中执行预处理
主进程只专注于进行预测
结果再次交给后台工作线程进行重采样和（可选的）导出

优点：

最适合预测大量图像
对 RAM 的使用更友好

缺点：

对于单个图像的预测不理想
需要图像以文件形式存在

    from nnunetv2.paths import nnUNet_results, nnUNet_raw
    import torch
    from batchgenerators.utilities.file_and_folder_operations import join
    from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
    
    # instantiate the nnUNetPredictor
    predictor = nnUNetPredictor(
        tile_step_size=0.5,
        use_gaussian=True,
        use_mirroring=True,
        perform_everything_on_device=True,
        device=torch.device('cuda', 0),
        verbose=False,
        verbose_preprocessing=False,
        allow_tqdm=True
    )
    # initializes the network architecture, loads the checkpoint
    predictor.initialize_from_trained_model_folder(
        join(nnUNet_results, 'Dataset003_Liver/nnUNetTrainer__nnUNetPlans__3d_lowres'),
        use_folds=(0,),
        checkpoint_name='checkpoint_final.pth',
    )
    # variant 1: give input and output folders
    predictor.predict_from_files(join(nnUNet_raw, 'Dataset003_Liver/imagesTs'),
                                 join(nnUNet_raw, 'Dataset003_Liver/imagesTs_predlowres'),
                                 save_probabilities=False, overwrite=False,
                                 num_processes_preprocessing=2, num_processes_segmentation_export=2,
                                 folder_with_segs_from_prev_stage=None, num_parts=1, part_id=0)

你也可以直接提供具体的文件，而不是输入和输出文件夹。如果提供具体文件，则不再需要 _0000 后缀！这在你无法控制文件名的情况下非常有用。记住，文件必须以“列表的列表”形式给出，其中外层列表的每个条目代表要预测的一个案例，内层列表包含属于该案例的所有文件。对于仅有一个输入模态的数据集（如 CT），内层列表只有一个文件，而对于其他模态（如有多个模态的 MRI，可能包括 T1、T2、Flair 等），内层列表可能有多个文件。
重要：每个案例的文件顺序必须与 dataset.json 中定义的通道顺序一致！
如果使用文件作为输入，你需要提供单独的输出文件作为输出！

    # variant 2, use list of files as inputs. Note how we use nested lists!!!
    indir = join(nnUNet_raw, 'Dataset003_Liver/imagesTs')
    outdir = join(nnUNet_raw, 'Dataset003_Liver/imagesTs_predlowres')
    predictor.predict_from_files([[join(indir, 'liver_152_0000.nii.gz')], 
                                  [join(indir, 'liver_142_0000.nii.gz')]],
                                 [join(outdir, 'liver_152.nii.gz'),
                                  join(outdir, 'liver_142.nii.gz')],
                                 save_probabilities=False, overwrite=False,
                                 num_processes_preprocessing=2, num_processes_segmentation_export=2,
                                 folder_with_segs_from_prev_stage=None, num_parts=1, part_id=0)

你知道吗？如果你不指定输出文件，预测的分割结果将直接返回：

    # variant 2.5, returns segmentations
    indir = join(nnUNet_raw, 'Dataset003_Liver/imagesTs')
    predicted_segmentations = predictor.predict_from_files([[join(indir, 'liver_152_0000.nii.gz')],
                                  [join(indir, 'liver_142_0000.nii.gz')]],
                                 None,
                                 save_probabilities=False, overwrite=True,
                                 num_processes_preprocessing=2, num_processes_segmentation_export=2,
                                 folder_with_segs_from_prev_stage=None, num_parts=1, part_id=0)

B.从 npy 数组进行预测

简而言之：

输入的图像是一个 npy 数组的列表
后台工作线程进行预处理
主进程仅专注于进行预测
结果再次交给后台线程进行重采样和（可选的）导出

优点：

适用于图像已经在内存中的情况
非常适合预测多张图像

缺点：

比默认模式占用更多的内存
不适合处理大量图像，因为所有图像都必须存放在内存中

from nnunetv2.imageio.simpleitk_reader_writer import SimpleITKIO

    img, props = SimpleITKIO().read_images([join(nnUNet_raw, 'Dataset003_Liver/imagesTs/liver_147_0000.nii.gz')])
    img2, props2 = SimpleITKIO().read_images([join(nnUNet_raw, 'Dataset003_Liver/imagesTs/liver_146_0000.nii.gz')])
    img3, props3 = SimpleITKIO().read_images([join(nnUNet_raw, 'Dataset003_Liver/imagesTs/liver_145_0000.nii.gz')])
    img4, props4 = SimpleITKIO().read_images([join(nnUNet_raw, 'Dataset003_Liver/imagesTs/liver_144_0000.nii.gz')])
    # we do not set output files so that the segmentations will be returned. You can of course also specify output
    # files instead (no return value on that case)
    ret = predictor.predict_from_list_of_npy_arrays([img, img2, img3, img4],
                                                    None,
                                                    [props, props2, props3, props4],
                                                    None, 2, save_probabilities=False,
                                                    num_processes_segmentation_export=2)

C.预测单个 npy 数组

简要说明：

输入为一个 npy 数组的图像
轴的顺序必须与对应的训练数据匹配。最简单的方法是使用与 nnU-Net 预处理时相同的 I/O 类加载图像。你可以在 plans.json 文件中的 image_reader_writer 键下找到该类。如果你选择自己处理，请注意医学图像的默认轴顺序是 SimpleITK 的顺序。如果使用 nibabel 加载图像，则需要将轴和间距从 [x,y,z] 转换为 [z,y,x]。
所有操作都在主进程中完成：预处理、预测、重采样、（导出）
无并行处理，是最慢的变体！
只有在无法一次性给 nnU-Net 提供多张图像时才使用此方法

优点：

无需处理多进程
无需处理数据迭代器等问题

缺点：

非常慢
除非你只能一次输入一张图像，否则永远不推荐使用这个选项

# predict a single numpy array (SimpleITKIO)
    img, props = SimpleITKIO().read_images([join(nnUNet_raw, 'Dataset003_Liver/imagesTr/liver_63_0000.nii.gz')])
    ret = predictor.predict_single_npy_array(img, props, None, None, False)

    # predict a single numpy array (NibabelIO)
    img, props = NibabelIO().read_images([join(nnUNet_raw, 'Dataset003_Liver/imagesTr/liver_63_0000.nii.gz')])
    ret = predictor.predict_single_npy_array(img, props, None, None, False)

    # The following IS NOT RECOMMENDED. Use nnunetv2.imageio!
    # nibabel, we need to transpose axes and spacing to match the training axes ordering for the nnU-Net default:
    nib.load('Dataset003_Liver/imagesTr/liver_63_0000.nii.gz')
    img = np.asanyarray(img_nii.dataobj).transpose([2, 1, 0])  # reverse axis order to match SITK
    props = {'spacing': img_nii.header.get_zooms()[::-1]}      # reverse axis order to match SITK
    ret = predictor.predict_single_npy_array(img, props, None, None, False)

D.使用自定义数据迭代器进行预测

简要说明：

非常灵活
不适合新手

优点：

你可以完全自行处理
拥有绝对的自由
如果记得在迭代器中使用多进程，会非常快

缺点：

所有事情都需要自己做
比你想象的要难

img, props = SimpleITKIO().read_images([join(nnUNet_raw, 'Dataset003_Liver/imagesTs/liver_147_0000.nii.gz')])
    img2, props2 = SimpleITKIO().read_images([join(nnUNet_raw, 'Dataset003_Liver/imagesTs/liver_146_0000.nii.gz')])
    img3, props3 = SimpleITKIO().read_images([join(nnUNet_raw, 'Dataset003_Liver/imagesTs/liver_145_0000.nii.gz')])
    img4, props4 = SimpleITKIO().read_images([join(nnUNet_raw, 'Dataset003_Liver/imagesTs/liver_144_0000.nii.gz')])
    # each element returned by data_iterator must be a dict with 'data', 'ofile' and 'data_properties' keys!
    # If 'ofile' is None, the result will be returned instead of written to a file
    # the iterator is responsible for performing the correct preprocessing!
    # note how the iterator here does not use multiprocessing -> preprocessing will be done in the main thread!
    # take a look at the default iterators for predict_from_files and predict_from_list_of_npy_arrays
    # (they both use predictor.predict_from_data_iterator) for inspiration!
    def my_iterator(list_of_input_arrs, list_of_input_props):
        preprocessor = predictor.configuration_manager.preprocessor_class(verbose=predictor.verbose)
        for a, p in zip(list_of_input_arrs, list_of_input_props):
            data, seg = preprocessor.run_case_npy(a,
                                                  None,
                                                  p,
                                                  predictor.plans_manager,
                                                  predictor.configuration_manager,
                                                  predictor.dataset_json)
            yield {'data': torch.from_numpy(data).contiguous().pin_memory(), 'data_properties': p, 'ofile': None}
    ret = predictor.predict_from_data_iterator(my_iterator([img, img2, img3, img4], [props, props2, props3, props4]),
                                               save_probabilities=False, num_processes_segmentation_export=3)

2.推理流程（A.直接从源文件预测，C.基于单个npy预测）

一共有四种预测方式，这里只给出这两种方式的预测流程。相对来说A难一点，C简单一点。
首先是通用的，实例化 nnUNetPredictor 类，并设置一些参数：

predictor = nnUNetPredictor(
    tile_step_size=0.5,
    use_gaussian=True,
    use_mirroring=True,
    perform_everything_on_device=True,
    device=torch.device('cuda', 0),
    verbose=False,
    verbose_preprocessing=False,
    allow_tqdm=True
)

具体细节见2.1 __init__ 方法
接下来，调用 initialize_from_trained_model_folder 方法来初始化模型：

predictor.initialize_from_trained_model_folder(
    join(nnUNet_results, 'Dataset003_Liver/nnUNetTrainer__nnUNetPlans__3d_lowres'),
    use_folds=(0,),
    checkpoint_name='checkpoint_final.pth',
)

具体细节见2.2 initialize_from_trained_model_folder 方法方法
另外，它在初始化模型的适合，实例化了plans_manager 这个类，用于配置模型的输入输出、数据预处理、网络架构等，确保推理时的设置与训练保持一致。

A.直接从源文件预测
…

C.基于单个npy预测
…

四.一些函数的具体实现细节以及知识点

滑动窗口

…

镜像翻转

…

高斯滤波

…