如何使用Keras对图像数据集进行预处理（数据增强、训练测试集拆分）

你原代码存在的核心问题 #

1. 基础语法错误：data_dir 赋值未加字符串引号、flow_from_directory 第一个参数传入了字符串"data_dir"而非变量、参数列表多处缺失逗号、路径拼接未做格式化处理。
2. 逻辑问题：flow_from_directory 的save_to_dir参数仅在迭代生成器、拉取批次数据时才会触发保存，直接定义生成器不会自动存储拆分后的图片；且当前ImageDataGenerator未配置任何数据增强参数，无法完成训练集增强需求；另外目标检测项目仅拆分图片的话，对应的标注文件没有同步匹配，会导致后续训练数据不对应。

完整实现方案 #

1. 环境依赖导入 #

import os
import shutil
from sklearn.model_selection import train_test_split
from PIL import Image
from tensorflow.keras.applications.xception import preprocess_input
from tensorflow.keras.preprocessing.image import ImageDataGenerator

2. 基础参数配置 #

# 按实际路径修改
original_data_dir = "/your/original/dataset/path"  # 原始数据集路径，下面按类别分子文件夹
output_root_dir = "/your/output/path"
train_dir = os.path.join(output_root_dir, "training")
test_dir = os.path.join(output_root_dir, "test")
augmented_train_dir = os.path.join(output_root_dir, "augmented_training")  # 增强后训练集存储路径

target_size = (150, 150)  # 统一调整的图像尺寸
test_split_ratio = 0.2  # 测试集占比
augment_per_image = 5  # 单张训练图生成的增强样本数量

3. 拆分数据集+统一调整尺寸并保存 #

# 先创建输出目录
os.makedirs(train_dir, exist_ok=True)
os.makedirs(test_dir, exist_ok=True)

# 遍历所有类别
for class_name in os.listdir(original_data_dir):
    class_path = os.path.join(original_data_dir, class_name)
    if not os.path.isdir(class_path):
        continue
    
    # 创建训练/测试集的类别子目录
    train_class_dir = os.path.join(train_dir, class_name)
    test_class_dir = os.path.join(test_dir, class_name)
    os.makedirs(train_class_dir, exist_ok=True)
    os.makedirs(test_class_dir, exist_ok=True)
    
    # 获取当前类别的所有图片路径
    all_images = [f for f in os.listdir(class_path) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    # 拆分训练测试集
    train_images, test_images = train_test_split(all_images, test_size=test_split_ratio, random_state=42)
    
    # 保存训练集图片（调整尺寸）
    for img_name in train_images:
        img = Image.open(os.path.join(class_path, img_name)).convert('RGB')
        img = img.resize(target_size)
        img.save(os.path.join(train_class_dir, img_name))
        # 如有标注文件，此处同步把对应标注文件复制到train_class_dir
    
    # 保存测试集图片（调整尺寸）
    for img_name in test_images:
        img = Image.open(os.path.join(class_path, img_name)).convert('RGB')
        img = img.resize(target_size)
        img.save(os.path.join(test_class_dir, img_name))
        # 同理同步复制对应标注文件到test_class_dir

4. 训练集数据增强实现 #

# 配置数据增强参数，可按需调整
datagen = ImageDataGenerator(
    preprocessing_function=preprocess_input,
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    horizontal_flip=True,
    zoom_range=0.2,
    shear_range=0.2,
    fill_mode='nearest'
)

# 遍历训练集每个类别做增强
for class_name in os.listdir(train_dir):
    class_path = os.path.join(train_dir, class_name)
    if not os.path.isdir(class_path):
        continue
    
    # 创建增强后存储目录
    aug_class_dir = os.path.join(augmented_train_dir, class_name)
    os.makedirs(aug_class_dir, exist_ok=True)
    
    # 生成增强样本
    aug_generator = datagen.flow_from_directory(
        train_dir,
        target_size=target_size,
        batch_size=1,
        classes=[class_name],
        save_to_dir=aug_class_dir,
        save_prefix='aug',
        save_format='jpg',
        shuffle=False
    )
    
    # 迭代生成器生成指定数量的增强样本
    total_aug_needed = len(os.listdir(class_path)) * augment_per_image
    count = 0
    for batch in aug_generator:
        count += 1
        if count >= total_aug_needed:
            break

关键注意点 #

• 目标检测项目拆分图片时必须同步拆分对应的标注文件，确保每张图片和其标注框的文件始终在同一个目录下，避免后续训练时数据不匹配。
• 数据增强的参数可根据项目场景调整，比如医疗图像不适合做水平翻转的话就关闭对应参数。
• 如果不需要额外保存增强后的图片到本地，也可以直接把增强生成器接入模型训练接口，不需要本地存储步骤，节省磁盘空间。

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