如何将imgaug数据增强应用于TensorFlow的dataset_from_directory数据集？

核心问题根源 #

tf.keras.utils.image_dataset_from_directory输出的是TensorFlow张量，而imgaug只支持numpy数组；直接用map会因为图模式下无法直接将张量转numpy报错，你遇到的占位符错误就是因为在图模式中强行执行numpy操作，导致TensorFlow无法追踪张量流程。

正确解决步骤 #

1. 编写兼容的增强函数

需要先把张量转numpy数组，用imgaug处理后再转回张量，同时注意图像数值范围的转换（imgaug更适合处理0-255的uint8格式）：

import numpy as np
import imgaug.augmenters as iaa
import tensorflow as tf

# 定义imgaug增强器
seq = iaa.Sequential([
    iaa.Fliplr(0.5),  # 水平翻转
    iaa.RandomBrightnessContrast(brightness=0.2, contrast=0.2),
    iaa.Affine(rotate=(-15, 15))  # 旋转
])

def augment_batch(images, labels):
    # 如果你的数据集已经做了归一化（0-1的float），先转回0-255的uint8
    images_np = (images.numpy() * 255).astype(np.uint8)
    # 执行增强
    augmented_imgs_np = seq(images=images_np)
    # 转回0-1的float张量
    augmented_imgs = tf.convert_to_tensor(augmented_imgs_np / 255.0, dtype=tf.float32)
    return augmented_imgs, labels

2. 用tf.py_function包装增强函数

因为imgaug是纯numpy操作，必须用tf.py_function把它包装成TensorFlow能识别的图节点，同时手动恢复张量形状（py_function会丢失形状信息）：

def tf_compatible_augment(images, labels):
    augmented_imgs, augmented_labels = tf.py_function(
        augment_batch,
        inp=[images, labels],
        Tout=[tf.float32, tf.int32]  # 要和你数据集的输出类型匹配，比如标签是int32
    )
    # 手动设置形状，避免后续模型报错
    augmented_imgs.set_shape(images.get_shape())
    augmented_labels.set_shape(labels.get_shape())
    return augmented_imgs, augmented_labels

3. 应用到训练数据集

# 你的原始数据集创建代码
train_ds = tf.keras.utils.image_dataset_from_directory(
    "你的训练集路径",
    image_size=(224, 224),
    batch_size=32,
    label_mode="int"
)

# 应用增强，加AUTOTUNE提升效率
train_ds_aug = train_ds.map(tf_compatible_augment, num_parallel_calls=tf.data.AUTOTUNE)

额外注意事项 #

• 如果你的数据集没有做归一化（输出是0-255的uint8），可以去掉代码中*255和/255.0的转换步骤；
• 如果遇到GPU兼容性问题，可以在增强函数里临时切换到CPU执行，比如在augment_batch开头加with tf.device("/cpu:0"):；
• 测试时可以先取一个batch验证增强效果，避免后续训练出问题。

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