TensorFlow图像增强后样本未有效扩充的问题求解

你的代码当前存在两个核心问题：

• 增强后的张量保留了(增强次数, batch_size, H, W, C)和(增强次数, batch_size, 1)的嵌套维度，导致每个原batch被处理成包含10个子batch的数据集元素，而非将增强样本平铺为独立样本。
• 使用map而非flat_map处理返回的子数据集，使得子数据集被当作单个元素，无法展开到主数据集中。

下面提供两种可行的解决方案：

方式一：基于原batch处理后展开 #

直接对每个batch的样本执行增强，合并嵌套维度后展开为独立样本：

batch_size = 16 
image_size = (400, 400)  # 补充缺失的图像尺寸定义

def augment_image(image):
    # 在这里添加你的图像增强操作，示例：
    # image = tf.image.random_flip_left_right(image)
    # image = tf.image.random_brightness(image, max_delta=0.2)
    return image

def augment_multiple_times(image, label, augment_count):
    # 对batch内所有图像执行augment_count次增强
    augmented_images = tf.stack([augment_image(image) for _ in range(augment_count)])
    # 合并增强次数与batch_size维度，得到(augment_count*batch_size, H, W, C)
    augmented_images = tf.reshape(augmented_images, (-1, image_size[0], image_size[1], 3))
    
    # 重复标签对应次数，合并维度得到(augment_count*batch_size, 1)
    augmented_labels = tf.repeat(label, augment_count, axis=0)
    
    print(f"Shape of augmented images: {augmented_images.shape}")
    print(f"Shape of labels: {augmented_labels.shape}")
    
    return augmented_images, augmented_labels

AUTOTUNE = tf.data.AUTOTUNE

train_dataset = tf.keras.utils.image_dataset_from_directory(
    input_path,
    image_size=image_size,
    batch_size=batch_size,
    label_mode="binary",
    validation_split=0.2,
    subset="training",
    seed=123,
)
print(f"Before processing: {train_dataset.cardinality().numpy()}")

# 使用flat_map替代map，将子数据集的元素平铺到主数据集
train_dataset = train_dataset.flat_map(
    lambda x, y: tf.data.Dataset.from_tensor_slices(augment_multiple_times(x, y, 10))
)

# 重新batch并打乱
train_dataset = train_dataset.batch(batch_size).shuffle(1000).prefetch(buffer_size=AUTOTUNE)
print(f"After processing: {train_dataset.cardinality().numpy()}")

方式二：先拆分为单个样本再增强 #

先将数据集拆分为单个样本，对每个样本增强后重新batch，逻辑更直观：

batch_size = 16 
image_size = (400, 400)

def augment_image(image):
    # 添加你的图像增强操作
    return image

def augment_single_sample(image, label, augment_count):
    # 对单个样本生成augment_count次增强结果
    augmented_images = tf.stack([augment_image(image) for _ in range(augment_count)])
    augmented_labels = tf.stack([label for _ in range(augment_count)])
    return augmented_images, augmented_labels

AUTOTUNE = tf.data.AUTOTUNE

train_dataset = tf.keras.utils.image_dataset_from_directory(
    input_path,
    image_size=image_size,
    batch_size=batch_size,
    label_mode="binary",
    validation_split=0.2,
    subset="training",
    seed=123,
)
print(f"Before processing: {train_dataset.cardinality().numpy()}")

# 先拆分为单个样本
train_dataset = train_dataset.unbatch()

# 对每个样本增强后平铺，再重新batch
train_dataset = train_dataset.flat_map(
    lambda x, y: tf.data.Dataset.from_tensor_slices(augment_single_sample(x, y, 10))
)

train_dataset = train_dataset.batch(batch_size).shuffle(1000).prefetch(buffer_size=AUTOTUNE)
print(f"After processing: {train_dataset.cardinality().numpy()}")

修改后效果 #

增强后的图像形状会变为(160, 400, 400, 3)（对应原16个样本×10次增强），标签形状变为(160, 1)，最终数据集的cardinality会变为原来的10倍（9×10=90个batch），每个元素都是独立的增强样本组成的batch，可直接用于模型训练。

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