如何对图像数据集执行Data Augmentation并保存至对应分类文件夹？

所需依赖 #

先安装必要的Python库：

pip install pillow torchvision

实现思路 #

1. 遍历原始数据集的每个分类文件夹（class1、class2）
2. 为每个分类在目标目录下创建对应的文件夹
3. 对每张原始图像应用多种数据增强变换
4. 将增强后的图像按类别保存到目标文件夹中

完整代码 #

import os
from PIL import Image
from torchvision import transforms

# 配置路径，根据你的实际文件夹路径修改
ORIGINAL_DIR = "./original_images"
AUGMENTED_DIR = "./augmented_images"

# 定义数据增强变换，可根据需求调整参数或增减变换
augment_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),  # 50%概率随机水平翻转
    transforms.RandomVerticalFlip(p=0.5),    # 50%概率随机垂直翻转
    transforms.RandomRotation(degrees=15),   # 随机旋转±15度
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),  # 随机调整色彩参数
])

# 创建增强数据集的根目录，已存在则跳过
os.makedirs(AUGMENTED_DIR, exist_ok=True)

# 遍历每个分类文件夹
for class_name in os.listdir(ORIGINAL_DIR):
    class_original_path = os.path.join(ORIGINAL_DIR, class_name)
    # 跳过非文件夹的文件
    if not os.path.isdir(class_original_path):
        continue
    
    # 创建对应分类的增强文件夹
    class_augmented_path = os.path.join(AUGMENTED_DIR, class_name)
    os.makedirs(class_augmented_path, exist_ok=True)
    
    # 遍历当前分类下的所有图像
    for img_name in os.listdir(class_original_path):
        img_path = os.path.join(class_original_path, img_name)
        # 只处理png格式，需要的话可添加jpg、jpeg等格式
        if not img_name.endswith(".png"):
            continue
        
        try:
            # 打开图像并统一转为RGB格式
            img = Image.open(img_path).convert("RGB")
            # 每张原图生成3张增强图，可修改数字调整生成数量
            for aug_idx in range(3):
                augmented_img = augment_transform(img)
                # 生成带序号的保存文件名，避免重名
                save_name = f"{os.path.splitext(img_name)[0]}_aug{aug_idx+1}.png"
                save_path = os.path.join(class_augmented_path, save_name)
                augmented_img.save(save_path)
            print(f"完成增强：{img_name}")
        except Exception as e:
            print(f"处理{img_name}出错：{str(e)}")

print("所有图像增强任务完成！")

关键说明 #

• 增强变换定制：可以修改augment_transform里的内容，比如添加transforms.RandomResizedCrop裁剪、调整旋转角度范围等，适配你的任务需求
• 生成数量调整：修改range(3)里的数字，就能改变每张原图生成的增强图像数量
• 图像格式兼容：若需要处理其他格式，把endswith(".png")改成endswith((".png", ".jpg", ".jpeg"))即可
• 路径适配：确保ORIGINAL_DIR和AUGMENTED_DIR的路径和你的实际文件夹路径一致

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