数据增强如何针对不同类别生成指定数量的新增样本？

原生ImageDataGenerator不直接支持按类别差异化配置增强规则、也不支持指定单个类别的新增样本数量，但是可以通过自定义逻辑实现需求，以下是两种常用实现方案：

方案1：离线生成指定类别的增强样本，合并到原数据集后训练 #

该方案可控性最高，可精确实现「给指定类别新增固定数量样本」的需求，适合小数据集场景：

• 第一步：从原数据集中筛选出需要增强的类别（示例需求中的类别0、1）的所有样本
• 第二步：初始化ImageDataGenerator配置增强参数，批量生成指定数量的增强样本，示例代码如下：

from tensorflow.keras.preprocessing.image import ImageDataGenerator
import pandas as pd
import numpy as np
import os

# 配置增强参数，可根据需求调整
datagen = ImageDataGenerator(
    rotation_range=15,
    width_shift_range=0.1,
    height_shift_range=0.1,
    horizontal_flip=True,
    fill_mode='nearest'
)

# 配置增强规则：目标类别、单类别新增样本数
target_classes = [0, 1]
aug_per_class = 1000
save_dir = './fer2013_augmented'
os.makedirs(save_dir, exist_ok=True)
new_sample_records = []

for cls in target_classes:
    # 筛选当前类别的所有原样本
    cls_original_samples = df[df['emotion'] == cls].reset_index(drop=True)
    generated_count = 0
    while generated_count < aug_per_class:
        # 随机抽取一张原样本作为增强基底
        base_sample = cls_original_samples.sample(n=1).iloc[0]
        # FER2013 csv存储的是像素值，此处转换为图像数组
        base_img = np.array(list(map(int, base_sample['pixels'].split()))).reshape(48, 48, 1)
        base_img = np.expand_dims(base_img, axis=0)
        
        # 生成增强样本并存储
        for aug_img in datagen.flow(
            base_img, 
            batch_size=1, 
            save_to_dir=save_dir, 
            save_prefix=f'aug_cls{cls}', 
            save_format='png'
        ):
            new_sample_records.append({
                'pixels': ' '.join(map(str, aug_img[0].astype(int).flatten())),
                'emotion': cls
            })
            generated_count += 1
            break

# 合并增强样本与原样本，得到最终训练数据集
aug_df = pd.DataFrame(new_sample_records)
final_train_df = pd.concat([df, aug_df], ignore_index=True)

• 第三步：直接使用合并后的final_train_df训练即可，无需在训练阶段额外配置增强。

方案2：训练阶段自定义生成器动态增强 #

如果不想占用额外磁盘空间存储离线增强样本，可以自定义继承自ImageDataGenerator的生成器，在训练取批次时仅对目标类别做动态增强，该方案适合不需要固定新增样本数量、仅想给少样本类别做加权增强的场景，示例代码如下：

from tensorflow.keras.preprocessing.image import ImageDataGenerator

class TargetClassAugGenerator(ImageDataGenerator):
    def __init__(self, target_classes, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.target_classes = target_classes
    
    def flow_from_dataframe(self, *args, **kwargs):
        base_generator = super().flow_from_dataframe(*args, **kwargs)
        while True:
            batch_x, batch_y = next(base_generator)
            # 仅对目标类别样本做增强
            for idx in range(len(batch_y)):
                if batch_y[idx] in self.target_classes:
                    batch_x[idx] = self.random_transform(batch_x[idx])
            yield batch_x, batch_y

# 生成器使用示例
train_generator = TargetClassAugGenerator(
    target_classes=[0,1],
    rotation_range=15,
    width_shift_range=0.1,
    height_shift_range=0.1,
    horizontal_flip=True
).flow_from_dataframe(
    dataframe=df,
    x_col='image_path', # 若使用像素值可自定义x_col读取逻辑
    y_col='emotion',
    target_size=(48,48),
    color_mode='grayscale',
    batch_size=32,
    class_mode='raw'
)

# 直接传入model.fit()训练即可
model.fit(train_generator, epochs=50)

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