nnU-Net 自定义图像 Reader/Writer 适配器完整指南：基于 BaseReaderWriter 扩展图像 I/O 支持

nnU-Net 通过 nnunetv2.imageio 模块中的 Reader/Writer 抽象层完成所有图像数据的读写，默认支持 NIfTI（nibabel / SimpleITK）、NRRD、MHA、PNG、BMP、TIF/TIFF 等格式。当你的数据集使用 nnU-Net 默认不支持的图像格式时，只需从 BaseReaderWriter 派生一个适配器并注册到读取器注册表，即可让 nnU-Net 的训练、预处理、推理与评估全流程透明地支持新格式。读完本文，你将掌握自定义图像 I/O 类的完整开发流程、三个抽象方法的契约细节、spacing 元数据的传递规则，以及如何通过 dataset.json 的 overwrite_image_reader_writer 字段让 nnU-Net 在运行时自动选用你的适配器。

一、为什么 nnU-Net 需要 Reader/Writer 抽象层

nnU-Net 的整个数据流水线——数据集指纹提取（fingerprint extraction）、预处理规划（planner）、训练（training）、推理（inference）、预测评估（evaluation）和结果可视化（overlay plots）——都依赖统一的图像读取接口。在 reader_writer_registry.py 中，determine_reader_writer_from_dataset_json 与 determine_reader_writer_from_file_ending 这两个函数负责在运行时根据 dataset.json 中的 file_ending 字段选择对应的 I/O 类。

目前仓库内置了五个开箱即用的实现（见 reader_writer_registry.py）：

类名	支持的文件后缀	说明
NaturalImage2DIO	.png, .bmp, .tif, .tiff	仅支持 2D 自然图像，基于 scikit-image 与 tifffile，见 natural_image_reader_writer.py
SimpleITKIO	.nii.gz, .nrrd, .mha, .gipl	基于 SimpleITK，同时记录 spacing / origin / direction 几何信息，见 simpleitk_reader_writer.py
Tiff3DIO	.tif, .tiff	3D TIFF 堆栈，支持用同名 .json 辅助文件提供 spacing 信息，见 tif_reader_writer.py
NibabelIO	.nii, .nii.gz	基于 nibabel，读取时做轴序转置与 SimpleITK 保持一致，见 nibabel_reader_writer.py
NibabelIOWithReorient	.nii, .nii.gz	在 NibabelIO 基础上将图像重定向到 RAS 标准方向，见 nibabel_reader_writer.py

当内置实现无法满足你的图像格式时，就需要编写自定义的 Reader/Writer。

二、自定义 Reader/Writer 的五个步骤

官方 readme.md 给出了从零编写适配器的完整清单，共五步：

1. 从 BaseReaderWriter 派生你的适配器；
2. 重新实现所有抽象方法；
3. 确保同时支持 2D 和 3D 输入图像（或者主动抛出明确错误）；
4. 将适配器文件放在 nnunetv2/imageio 目录下，否则 nnU-Net 找不到它；
5. 把它添加到 reader_writer_registry.py 的 LIST_OF_IO_CLASSES 列表中。

完成以上步骤后，你的适配器即可被 nnU-Net 自动发现并使用。第 5 步的注册至关重要：determine_reader_writer_from_file_ending 会遍历 LIST_OF_IO_CLASSES，比对 file_ending.lower() in rw.supported_file_endings 来挑选匹配的类，未注册的类永远不会被选中。

三、三个抽象方法：读写契约详解

BaseReaderWriter 定义在 base_reader_writer.py 中，继承自 Python 标准库的 ABC，包含三个必须实现的抽象方法。这是自定义适配器的核心契约，nnU-Net 的规划、预处理和推理逻辑都严格依赖以下返回约定。

1. read_images(image_fnames) -> (np.ndarray, dict)

读取一个或多个图像文件，返回一个四维 NumPy 数组与一个元数据字典：

• 数组形状必须是 (c, x, y, z)：第一个轴是模态/通道数（可以是 1），随后是空间维度。多模态数据的多个文件会通过 np.vstack 沿第一个轴堆叠（可参考 simpleitk_reader_writer.py 的 np.vstack(images, dtype=np.float32, casting='unsafe') 用法）。
• 2D 图像的特殊约定：返回数组形状应为 (c, 1, x, y)，且 spacing 必须设为 (999, sp_x, sp_y)，其中 999 必须大于 sp_x 和 sp_y。例如 RGB 图返回 shape=(3, 1, 224, 224), spacing=(999, 1, 1)（参考 base_reader_writer.py）。
• 字典用于存放转成 NumPy 数组后会丢失的元信息（如 Spacing、Orientation、Direction）。该字典会在推理时被原样交给 write_seg 用于导出预测分割，所以凡是导出时需要的元信息都必须放进去。
• 强制要求：字典必须包含 'spacing' 键，值为长度为 3 的 tuple/list，顺序必须与返回数组的轴顺序一一对应：数组形状为 c,x,y,z、spacing 为 (a,b,c)，则 a 对应 x 轴、b 对应 y 轴、c 对应 z 轴（见 base_reader_writer.py）。spacing 是后续规划与预处理的关键输入。
• 对没有 spacing 概念的图像，统一设为 1（2D 时首轴 999 的约定仍然适用）。

2. read_seg(seg_fname) -> (np.ndarray, dict)

与 read_images 的要求相同，但返回的分割数组形状必须为 (1, x, y, z)——nnU-Net 目前不允许单文件包含多个分割图。如果图像和分割可以用同一套逻辑读取，直接复用即可，如 simpleitk_reader_writer.py 的写法：

def read_seg(self, seg_fname: str) -> Tuple[np.ndarray, dict]:
    return self.read_images((seg_fname, ))

3. write_seg(seg, output_fname, properties) -> None

将预测分割导出为指定格式：

• 传入的 seg 数组与对应图像数据具有相同的形状与方向，无需自行重采样或调整方向，直接保存即可。
• 重要约定：分割始终是 3D 的！如果输入图像是 2D，则分割形状为 (1, x, y)，导出时需要先 seg = seg<a href="https://link.gitcode.com/i/61253d6ff152162d9f424bb65fc2c35d" target="_blank">0] 转回 2D（见 [base_reader_writer.py）。
• properties 就是 read_images / read_seg 阶段创建的那个字典，用于恢复元数据。
• 数据类型按标签最大值选择：np.max(seg) < 255 时用 uint8，否则用 uint16（参见 simpleitk_reader_writer.py 等实现）。

四、从内置实现中学习：可复用的关键模式

在动手编写自定义适配器之前，先看两个内置实现的经典做法，它们几乎覆盖了所有自定义场景会遇到的坑。

模式一：几何信息与 spacing 的顺序修正（SimpleITKIO）

SimpleITKIO 演示了最完整的元数据处理：

• sitk.GetArrayFromImage 返回的数组是 z,y,x 顺序，而 nnU-Net 要求 x,y,z，因此 spacing 必须用 <a href="https://link.gitcode.com/i/fde70c6d66b7c6de24b98d3b64d86311" target="_blank">::-1] 反转；注释里也直言这是 SimpleITK 轴序带来的常见坑（见 [simpleitk_reader_writer.py）。
• 2D 图像处理：npy_image<a href="https://link.gitcode.com/i/267e4e00bf57a870740d1b5c7407ef14" target="_blank">None, None] 补出通道轴与 x 轴，spacing 取 (max_spacing * 999, *list(spacings[-1])[::-1])，即用最大 spacing 的 999 倍占位首轴（见 [simpleitk_reader_writer.py）。
• 4D 多模态文件（单文件多模态）会被识别并取前三维的 spacing（见 simpleitk_reader_writer.py）。
• 读取后会做一致性检查：多模态文件形状、spacing 必须一致（不一致报 ERROR），origin、direction 不一致仅发 WARNING，并提示用 nnUNetv2_plot_overlay_pngs 验证数据与分割的重叠（见 simpleitk_reader_writer.py）。这些检查通过基类的 _check_all_same 与 _check_all_same_array 静态方法完成（base_reader_writer.py）。
• 几何信息存入 'sitk_stuff' 子字典，spacing 单独放 'spacing' 键；write_seg 时用 sitk.GetImageFromArray 重建并 SetSpacing/SetOrigin/SetDirection 恢复几何（见 simpleitk_reader_writer.py）。

模式二：方向重定向与轴序转换（NibabelIOWithReorient）

NibabelIOWithReorient 展示了"读入时统一方向、导出时还原方向"的完整闭环：

• 读取时用 io_orientation + as_reoriented 将图像重定向到 RAS；同时保存 original_affine 和 reoriented_affine 两个仿射矩阵（见 nibabel_reader_writer.py）。
• nibabel 的轴序与 SimpleITK 不同，需用 transpose((2, 1, 0)) 转置统一；注释明确提醒"IMPORTANT: Run nnUNetv2_plot_overlay_pngs to verify that this did not destroy the alignment of data and seg!"（见 nibabel_reader_writer.py）。
• 导出时先还原转置，再通过 ornt_transform 将 RAS 方向变换回原始方向，并校验恢复后的仿射与原始仿射是否一致，不一致时打印 WARNING（见 nibabel_reader_writer.py）。

这两个模式提醒我们：自定义适配器最需要关注的就是轴序约定与spacing 语义，读入与写出必须严格对称，否则数据与分割会产生错位。

五、Tiff3DIO：一个轻量自定义范本

Tiff3DIO 是仓库中"非医学格式适配器"的范本，值得重点参考：

• 仅支持 3D 图像：tifffile.imread 读入后检查 image.ndim != 3 即抛出 RuntimeError（见 tif_reader_writer.py），这正是文档要求"支持 2D 和 3D（或抛出错误）"的体现——2D TIFF 请使用 NaturalImage2DIO。
• spacing 缺失的兜底：TIFF 本身难以可靠写入分辨率信息，因此 Tiff3DIO 采用同名 JSON 辅助文件方案：图像 image1_0000.tif 对应的 spacing 文件应为 image1.json（省略通道标识符）；若找不到辅助文件，打印 WARNING 并默认 spacing 为 (1, 1, 1)（见 tif_reader_writer.py）。
• 写出时反向生成辅助文件：write_seg 用 zlib 压缩写出 3D TIFF，同时在同目录写出 <名字>.json 保存 spacing（见 tif_reader_writer.py），保证下一轮读取时几何信息不丢失。
• 分割的 spacing 校验：read_seg 中额外断言 all(<a href="https://link.gitcode.com/i/15e69f0836e120774ae1b2a9a8cdfdbe" target="_blank">i > 0 for i in spacing])（见 [tif_reader_writer.py）。

如果你的数据集格式与 TIFF 类似（无内嵌几何元数据），可以完全照搬这个"JSON 伴生文件"方案。

六、让 nnU-Net 运行时选用你的适配器

自定义类编写完成后，有两条路径让 nnU-Net 找到它。

路径一：默认自动匹配（无需改代码）

只要满足第 5 步（加入 LIST_OF_IO_CLASSES），且你的类定义了正确的 supported_file_endings 类属性，determine_reader_writer_from_file_ending 就会根据 dataset.json 的 file_ending 自动匹配。匹配成功后，注册表还会尝试用 example_file 实际读一次文件来验证可用性，读取失败则自动跳过该类并打印 traceback（见 reader_writer_registry.py）。

路径二：通过 dataset.json 显式指定（推荐）

在数据集的 dataset.json 中加入 overwrite_image_reader_writer 字段，值为你的类名（字符串）。determine_reader_writer_from_dataset_json 会优先检查该字段（且要求其值不为 'None' 字符串），并通过 recursive_find_reader_writer_by_name 在 nnunetv2.imageio 模块内递归查找同名类；找不到时打印警告并回退到按文件后缀自动匹配（见 reader_writer_registry.py、reader_writer_registry.py）。

生成 dataset.json 的官方入口是 generate_dataset_json.py，其 docstring 明确说明："If you need a special IO class for your dataset you can derive it from BaseReaderWriter, place it into nnunet.imageio and reference it here by name"。函数签名接受 overwrite_image_reader_writer 参数，并在第 104-105 行将该值写入 dataset_json['overwrite_image_reader_writer']。

仓库内的数据集转换脚本给出了实战用法，例如 Dataset220_KiTS2023.py、Dataset218_Amos2022_task1.py、Dataset224_AbdomenAtlas1.0.py 都通过 overwrite_image_reader_writer='NibabelIOWithReorient' 为数据量较大、方向不一的 CT/MRI 数据集强制启用方向重定向。

七、从源码结构看：适配器在整个流水线中的调用位置

自定义适配器并非只影响推理导出，它贯穿 nnU-Net 的完整生命周期，以下调用点均可在仓库源码中验证：

• 数据集完整性校验：verify_dataset_integrity.py 通过 determine_reader_writer_from_dataset_json 获取 I/O 类并尝试读取首个样本。
• 数据集指纹提取：fingerprint_extractor.py 在计算 shape、spacing 等统计量时使用该 I/O 类。
• 实验规划器：default_experiment_planner.py 在生成 plans 文件时调用。
• 预测评估：evaluate_predictions.py 同时使用 determine_reader_writer_from_dataset_json 与 determine_reader_writer_from_file_ending 读取预测与金标准。
• 结果可视化：overlay_plots.py 在绘制数据与分割叠加图时使用。
• 预训练权重迁移：move_plans_between_datasets.py 跨数据集移动 plans 时也会读取目标数据集的图像。

因此，只要适配器符合契约，训练、验证、推理、评估整条链路的读写都会自动切换到你的实现，无需修改任何其他代码。

八、实践清单与常见陷阱

开发清单

1. 新建文件放入 nnunetv2/imageio/ 目录（nnU-Net 只会在这个目录里找 I/O 类，放到别处将无法被发现）；
2. 类继承 BaseReaderWriter 并定义 supported_file_endings 列表（注意大小写不敏感匹配：determine_reader_writer_from_file_ending 使用 file_ending.lower() 比较）；
3. 实现 read_images、read_seg、write_seg 三个抽象方法；
4. 返回数组遵守 (c, x, y, z) 四维约定（2D 为 (c, 1, x, y)）；
5. 元数据字典必须含 'spacing' 键且顺序与数组轴一致（2D 首轴用 999）；
6. 将类加入 reader_writer_registry.py 的 LIST_OF_IO_CLASSES；
7. 在 dataset.json 中设置 overwrite_image_reader_writer 以显式指定（可选但推荐）；
8. 用 nnUNetv2_plot_overlay_pngs 验证数据与分割是否对齐。

常见陷阱

• 轴序不对称：读入时做了转置/反转，写出时必须对称还原，否则分割错位（NibabelIO 的 transpose((2,1,0)) 与写出的对称还原就是标准范式）；
• spacing 顺序错误：这是最常见且最隐蔽的问题——nnU-Net 的预处理重采样直接依赖 spacing，顺序与轴不对应会导致所有重采样结果错误；
• 2D 的 999 约定：2D 图像首轴 spacing 必须设为大于其余轴的 999，否则规划器无法正确识别 2D 数据；
• 忽略 dtype 上限：标签值超过 255 时仍用 uint8 会溢出，应遵循内置实现"np.max(seg) < 255 用 uint8、否则 uint16"的规则；
• 文件格式与分割格式不一致：dataset.json 的 file_ending 必须同时匹配图像与分割文件的后缀（generate_dataset_json.py 明确强调这一点）；
• 有损压缩不可用于标签：NaturalImage2DIO 特意不支持 .jpg/.jpeg，因为不允许有损压缩破坏分割图（见 natural_image_reader_writer.py 的注释）。

遵循以上五步与契约细节，再对照 SimpleITKIO、NibabelIOWithReorient、Tiff3DIO 三个内置范本，即可为任意图像格式编写出被 nnU-Net 全流程支持的 Reader/Writer 适配器。更深入的格式说明可继续阅读 imageio 模块文档 与数据集格式参考 dataset-format.md。