使用Keras ImageDataGenerator增强图像时如何保留原文件夹结构

这个问题我之前帮不少人解决过——核心就是复刻原始数据集的目录结构到增强目标路径，再针对性地把增强图存到对应子文件夹。下面给你两种实用的Python实现方案，直接就能套用到你的10类图像数据集上：

方案1：用PIL+os模块（基础通用） #

如果你只是需要简单的增强操作（翻转、旋转、亮度调整等），用Python自带的os模块加PIL库就能搞定，代码逻辑清晰，容易调试：

步骤拆解 #

• 遍历原始数据集的所有目录和文件
• 对每个子目录，在增强根路径下创建完全相同的文件夹结构
• 对每张图像，执行增强操作，保存到对应子文件夹（给增强图加后缀区分原文件）

完整代码示例 #

import os
from PIL import Image, ImageEnhance, ImageOps
import random

# 配置路径
original_root = "./your_original_dataset"  # 你的原始数据集根目录
augmented_root = "./augmented_dataset"     # 增强后保存的根目录
augment_count = 3  # 每张图生成3张增强图

# 定义增强函数
def augment_image(img):
    # 随机水平翻转
    if random.random() > 0.5:
        img = ImageOps.mirror(img)
    # 随机垂直翻转
    if random.random() > 0.5:
        img = ImageOps.flip(img)
    # 随机亮度调整
    enhancer = ImageEnhance.Brightness(img)
    img = enhancer.enhance(random.uniform(0.8, 1.2))
    # 随机旋转（-15到15度）
    img = img.rotate(random.randint(-15, 15), expand=True)
    return img

# 遍历原始目录
for root, dirs, files in os.walk(original_root):
    # 计算当前目录相对于原始根的路径，用来复刻结构
    relative_path = os.path.relpath(root, original_root)
    # 构建增强后的目标目录
    target_dir = os.path.join(augmented_root, relative_path)
    # 创建目录（如果不存在的话）
    os.makedirs(target_dir, exist_ok=True)
    
    # 处理每个文件
    for file in files:
        # 只处理图像文件（可根据需要添加更多格式）
        if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
            file_path = os.path.join(root, file)
            # 打开原始图像
            try:
                img = Image.open(file_path).convert('RGB')
                # 保存原始图像（可选，如果你想把原图也放进增强数据集的话）
                img.save(os.path.join(target_dir, file))
                # 生成增强图像
                for i in range(augment_count):
                    aug_img = augment_image(img.copy())
                    # 构建增强后的文件名，比如cat.jpg → cat_aug1.jpg
                    name, ext = os.path.splitext(file)
                    aug_file_name = f"{name}_aug{i+1}{ext}"
                    aug_file_path = os.path.join(target_dir, aug_file_name)
                    aug_img.save(aug_file_path)
            except Exception as e:
                print(f"处理文件{file_path}出错: {e}")
                continue

方案2：用Albumentations库（专业增强，支持更多变换） #

如果需要更复杂的增强（比如随机裁剪、高斯模糊、颜色抖动等），推荐用Albumentations库，它是计算机视觉领域常用的高效增强工具：

先安装库 #

pip install albumentations

完整代码示例 #

import os
import cv2
import albumentations as A

# 配置路径
original_root = "./your_original_dataset"
augmented_root = "./augmented_dataset"
augment_count = 3

# 定义增强管道（可根据需求自定义变换）
transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.VerticalFlip(p=0.5),
    A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.7),
    A.RandomRotate90(p=0.3),
    A.GaussianBlur(blur_limit=(3, 7), p=0.2),
    # 最后转成RGB（避免灰度图问题）
    A.ToRGB(),
])

# 遍历目录并处理
for root, dirs, files in os.walk(original_root):
    relative_path = os.path.relpath(root, original_root)
    target_dir = os.path.join(augmented_root, relative_path)
    os.makedirs(target_dir, exist_ok=True)
    
    for file in files:
        if file.lower().endswith(('.png', '.jpg', '.jpeg')):
            file_path = os.path.join(root, file)
            # 用OpenCV读取图像（Albumentations默认用BGR格式，最后会转成RGB）
            img = cv2.imread(file_path)
            if img is None:
                print(f"无法读取文件{file_path}")
                continue
            # 保存原图（可选）
            cv2.imwrite(os.path.join(target_dir, file), img)
            # 生成增强图
            for i in range(augment_count):
                augmented = transform(image=img)
                aug_img = augmented['image']
                name, ext = os.path.splitext(file)
                aug_file_name = f"{name}_aug{i+1}{ext}"
                aug_file_path = os.path.join(target_dir, aug_file_name)
                cv2.imwrite(aug_file_path, aug_img)

关键小贴士 #

1. 文件名区分：给增强图加_augN后缀，避免和原图或其他增强图重名
2. 可选保存原图：如果你的训练需要同时用原图和增强图，就保留代码里保存原图的部分；如果只需要增强图，删掉那行就行
3. 多进程加速：如果数据集很大，可结合multiprocessing模块并行处理，提升速度
4. 格式兼容：代码里只处理了常见图像格式，可根据你的数据集添加更多格式（比如.tiff）

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