TensorFlow图像分类数据增强未增大数据集大小的问题咨询

原因分析 #

• 你使用的RandomFlip/RandomRotation等增强属于动态实时增强：这类操作是在模型训练时，每次读取原样本后临时生成增强图像，不会提前生成并扩充数据集的样本总数，因此train_ds的长度始终和原数据集一致，不会自动变为4倍规模。
• 直接合并原数据集与增强后的数据集，本质是将原样本和它们的单次增强结果合并，所以只会得到2倍规模，无法实现多次增强后的多倍扩充效果。

解决办法 #

要实现数据集的物理扩充（生成并保存增强图像到文件夹），可以通过批量生成增强图像并按类别保存的方式完成，具体步骤如下：

1. 准备目录结构 #

假设原训练数据的目录结构为：

train_dir/
    class_1/
        img1.jpg
        img2.jpg
    class_2/
        img1.jpg
        ...

创建一个新目录用于保存增强后的图像，比如augmented_train_dir，并为每个类别创建对应的子文件夹。

2. 批量生成并保存增强图像 #

使用以下代码实现增强图像的生成与保存（每个原样本生成3份增强图，最终数据集规模为原数据的4倍）：

import tensorflow as tf
import os
from PIL import Image

# 替换为你的原训练数据目录和增强数据保存目录
train_dir = "path/to/your/train_dir"
augmented_train_dir = "path/to/augmented_train_dir"

# 定义数据增强序列
data_augmentation = tf.keras.Sequential([
    tf.keras.layers.RandomFlip("horizontal"),
    tf.keras.layers.RandomRotation(0.2),
    tf.keras.layers.RandomZoom(0.2)
])

# 加载原训练数据集（单样本批量处理，方便获取文件名和标签）
train_ds = tf.keras.preprocessing.image_dataset_from_directory(
    train_dir,
    image_size=(224, 224),  # 替换为你的图像尺寸
    batch_size=1,
    label_mode='int'
)

# 创建增强数据的类别文件夹
class_names = train_ds.class_names
for class_name in class_names:
    os.makedirs(os.path.join(augmented_train_dir, class_name), exist_ok=True)

# 每个原样本生成3份增强图，可按需调整数量
augment_count_per_sample = 3

# 遍历数据集生成并保存图像
for idx, (image, label) in enumerate(train_ds):
    # 获取原图像的文件名和后缀
    original_path = train_ds.file_paths[idx]
    file_name = os.path.basename(original_path).split('.')[0]
    file_ext = os.path.basename(original_path).split('.')[1]
    current_class = class_names[label[0]]
    
    # 保存原图像到增强目录（可选，若需要合并原数据与增强数据）
    original_img = tf.keras.preprocessing.image.array_to_img(image[0])
    original_img.save(os.path.join(augmented_train_dir, current_class, f"{file_name}_original.{file_ext}"))
    
    # 生成并保存增强图像
    for i in range(augment_count_per_sample):
        aug_image = data_augmentation(image, training=True)
        aug_img = tf.keras.preprocessing.image.array_to_img(aug_image[0])
        aug_img.save(os.path.join(augmented_train_dir, current_class, f"{file_name}_aug_{i+1}.{file_ext}"))

3. 加载扩充后的数据集 #

完成增强图像保存后，直接用image_dataset_from_directory加载新的数据集即可：

augmented_train_ds = tf.keras.preprocessing.image_dataset_from_directory(
    augmented_train_dir,
    image_size=(224, 224),
    batch_size=32,
    label_mode='int'
)

注意事项 #

• 增强操作的随机性确保每次生成的图像不同，同一个样本多次增强会得到不同结果。
• 若原数据集规模较大，生成增强图像会占用较多磁盘空间，可根据实际需求调整augment_count_per_sample的数值。

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