如何在TensorFlow中构建适配3D体数据的分割任务数据管道

结论 #

完全可以沿用你现有的2D分割tf.data流水线结构，仅需要替换2D图像解码逻辑为3D数据适配的读取解码逻辑，最终可输出形状为(BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, IMG_DEPTH, CHANNELS)的批次，且不会一次性加载全量数据到内存。

实现方案 #

方案1：原生TensorFlow API实现（无额外依赖） #

你可以提前将所有3D图像、mask导出为固定维度的.raw原始二进制文件，单通道uint8格式下每个样本的文件大小固定为IMG_HEIGHT * IMG_WIDTH * IMG_DEPTH字节，直接在map阶段读取后reshape即可。
示例代码如下：

import tensorflow as tf

# 提前定义3D数据的固定维度，可根据你的需求调整
IMG_HEIGHT = 128
IMG_WIDTH = 128
IMG_DEPTH = 64
CHANNELS = 1
BATCH_SIZE = 4

@tf.function
def parse_3d_function(path_pair):
    img_path = path_pair[0]
    mask_path = path_pair[1]
    
    # 读取3D图像的二进制文件
    img_raw = tf.io.read_file(img_path)
    img = tf.io.decode_raw(img_raw, tf.uint8)
    img = tf.reshape(img, (IMG_HEIGHT, IMG_WIDTH, IMG_DEPTH, CHANNELS))
    
    # 同理处理对应mask
    mask_raw = tf.io.read_file(mask_path)
    mask = tf.io.decode_raw(mask_raw, tf.uint8)
    mask = tf.reshape(mask, (IMG_HEIGHT, IMG_WIDTH, IMG_DEPTH, CHANNELS))
    
    # 可在此处添加3D数据增强逻辑，如随机翻转、随机裁剪等
    return img, mask

# 流水线后续逻辑和你2D场景的用法完全一致
path_pairs = [] # 替换为你自己的(图像路径, mask路径)元组列表
dataset = tf.data.Dataset.from_tensor_slices(path_pairs)
dataset = dataset.shuffle(len(path_pairs))
# 开启并行映射加速解码
dataset = dataset.map(parse_3d_function, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.batch(BATCH_SIZE)
dataset = dataset.prefetch(tf.data.AUTOTUNE)

方案2：适配NIfTI格式3D数据 #

如果你使用的是医学影像常用的NIfTI（.nii/.nii.gz）格式，可安装对应版本的tensorflow-io扩展包，使用tfio.image.decode_nifti接口直接解码3D文件，解码后补充通道维度即可使用，无需提前转.raw格式。

注意事项 #

• 若你的原始3D数据尺寸不统一，可在parse_3d_function中添加裁剪/填充逻辑，将所有样本对齐到相同的固定维度后再返回，才能正常生成批次
• 若需要更高的读取性能，可提前将3D数据写入TFRecord格式存储，流式读取效率会高于散存的单文件

内容的提问来源于stack exchange，提问作者user497663