如何为含图像与掩码的tf.data.Dataset执行同步数据增强？

针对你基于tf.data.Dataset构建的图像-掩码数据集，我们可以通过自定义同步增强函数实现两者的一致变换，全程使用TensorFlow原生API，完全避开已废弃的ImageDataGenerator。

核心逻辑 #

让图像和掩码共享同一随机状态，确保每一步增强的参数（如翻转方向、旋转角度、裁剪区域）完全一致；同时针对掩码的整数标签属性，选择合适的插值方式避免标签失真。

具体实现代码 #

1. 定义同步增强函数

import tensorflow as tf

def augment_image_mask(image, mask):
    # 图像转浮点型（适配后续操作），掩码保留原类型
    image = tf.cast(image, tf.float32)
    
    # 随机水平翻转（共用随机判断结果）
    if tf.random.uniform(()) > 0.5:
        image = tf.image.flip_left_right(image)
        mask = tf.image.flip_left_right(mask)
    
    # 随机垂直翻转
    if tf.random.uniform(()) > 0.5:
        image = tf.image.flip_up_down(image)
        mask = tf.image.flip_up_down(mask)
    
    # 随机旋转0/90/180/270度（共用随机旋转次数）
    rotate_k = tf.random.uniform(shape=[], minval=0, maxval=4, dtype=tf.int32)
    image = tf.image.rot90(image, k=rotate_k)
    mask = tf.image.rot90(mask, k=rotate_k)
    
    # 随机裁剪（合并图像和掩码后同步裁剪，保证区域一致）
    combined = tf.concat([image, mask], axis=-1)
    cropped_combined = tf.image.random_crop(combined, size=[image.shape[0], image.shape[1], image.shape[2]+mask.shape[2]])
    image = cropped_combined[..., :image.shape[2]]
    mask = cropped_combined[..., image.shape[2]:]
    
    # 图像归一化（可选，根据你的任务需求调整）
    image = image / 255.0
    
    # 确保掩码回到原整数类型（比如uint8）
    mask = tf.cast(mask, tf.uint8)
    
    return image, mask

2. 将增强函数映射到数据集

# 应用增强操作，用AUTOTUNE自动并行提升处理效率
val_data = val_data.map(augment_image_mask, num_parallel_calls=tf.data.AUTOTUNE)

关键细节说明 #

• 同步随机控制：所有随机判断和参数只生成一次，同时作用于图像和掩码，保证变换完全同步。
• 掩码保护：旋转、裁剪操作对掩码默认使用nearest（最近邻）插值，避免整数标签出现无效的浮点值。
• 性能优化：num_parallel_calls=tf.data.AUTOTUNE让TensorFlow自动调度并行处理，加快数据加载和增强速度。

扩展自定义操作 #

如果需要仅针对图像的增强（如亮度、对比度调整），直接单独处理图像即可，掩码无需变动：

# 示例：仅对图像做随机亮度调整
image = tf.image.random_brightness(image, max_delta=0.1)
image = tf.clip_by_value(image, 0.0, 1.0)  # 限制像素值在合理范围

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