如何使用Keras实现BreakHis数据集的指定规则数据增强

需求说明 #

以BreakHis数据集中的DC子类样本量为基准，对其余所有子类的图像分别执行上下翻转、左右翻转、90°逆时针旋转、180°逆时针旋转4种固定变换完成扩增，实现类别样本量均衡。

原始数据集分布：

增强后数据集分布：

Keras实现代码 #

方案1：离线扩增（生成扩增后图片存本地，适配任何训练流程） #

import tensorflow as tf
from tensorflow.keras import layers
import os
import shutil

# 扩增参数配置
CLASS_TO_KEEP = "DC"  # 基准类，不做扩增
AUG_METHODS = [
    ("vflip", layers.RandomFlip("vertical")),
    ("hflip", layers.RandomFlip("horizontal")),
    ("rot90", layers.RandomRotation(factor=(-0.25, -0.25), fill_mode="reflect")),
    ("rot180", layers.RandomRotation(factor=(-0.5, -0.5), fill_mode="reflect"))
]
RAW_DATASET_PATH = "替换为你的BreakHis原始数据集路径"
AUG_OUTPUT_PATH = "替换为扩增后数据集存储路径"

def augment_single_class(raw_class_path, output_class_path):
    # 先复制原始图片到输出目录
    shutil.copytree(raw_class_path, output_class_path)
    # 遍历所有原始图片生成扩增图
    for img_name in os.listdir(raw_class_path):
        if not img_name.lower().endswith((".png", ".jpg", ".jpeg")):
            continue
        # 读取原图
        raw_img = tf.io.read_file(os.path.join(raw_class_path, img_name))
        raw_img = tf.io.decode_image(raw_img, channels=3)
        # 依次应用4种变换，各生成1张扩增图
        for prefix, aug_layer in AUG_METHODS:
            aug_img = aug_layer(raw_img, training=True)
            tf.keras.utils.save_img(
                os.path.join(output_class_path, f"{prefix}_{img_name}"),
                aug_img.numpy()
            )

if __name__ == "__main__":
    os.makedirs(AUG_OUTPUT_PATH, exist_ok=True)
    # 基准类直接复制，不扩增
    shutil.copytree(
        os.path.join(RAW_DATASET_PATH, CLASS_TO_KEEP),
        os.path.join(AUG_OUTPUT_PATH, CLASS_TO_KEEP)
    )
    # 其余所有类执行扩增
    for class_name in os.listdir(RAW_DATASET_PATH):
        raw_class_path = os.path.join(RAW_DATASET_PATH, class_name)
        if class_name == CLASS_TO_KEEP or not os.path.isdir(raw_class_path):
            continue
        augment_single_class(raw_class_path, os.path.join(AUG_OUTPUT_PATH, class_name))

方案2：在线扩增（训练时实时生成扩增图，无需提前存本地） #

如果不需要保留扩增后的图片，可以直接在训练流程中嵌入扩增层，仅对非DC类样本应用变换：

import tensorflow as tf
from tensorflow.keras import layers

# 构建扩增层
aug_layer = tf.keras.Sequential([
    layers.RandomFlip("vertical_and_horizontal"),
    layers.RandomRotation(factor=(-0.5, -0.25), fill_mode="reflect")
])

# 自定义预处理函数，仅对非DC类样本做扩增
def preprocess_fn(image, label, dc_label_id=0):
    # label为DC类对应的标签id时不做扩增，其余类别做扩增
    if label != dc_label_id:
        image = aug_layer(image, training=True)
    return image, label

# 加载数据集时应用预处理函数即可
train_dataset = tf.keras.utils.image_dataset_from_directory(
    RAW_DATASET_PATH,
    image_size=(224, 224),
    batch_size=32
).map(preprocess_fn)

通用实现方法 #

这类指定类别固定变换的均衡扩增，通用逻辑如下：

• 先统计所有类别的样本量，确定基准类的样本数，计算其余类别需要的扩增倍数，匹配基准类的样本量级
• 固定变换规则，不要使用随机扩增，保证每张原图生成固定数量的扩增样本，避免样本量不可控
• 扩增样本和原图归为同一类别，不修改原有标签映射关系
• 若需要划分训练/验证/测试集，必须先完成数据集拆分，再仅对训练集做扩增，禁止先扩增再拆分，避免数据泄漏导致模型泛化性下降

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