教你使用Python实现遍历文件夹进行图像增强

1. 问题背景

当做图像分类等计算机视觉任务时，数据集中图像数量及质量，直接影响着深度学习模型训练的效果。对数据集图像进行增强扩充是一个很好的解决办法。

2. 环境要求

• Python
• Opencv-python
• imgaug

3. 实现代码

如： 这是待增强的图像文件夹，一共四个类别，每个类别文件夹里包含着相应的图像。

datasets 文件夹同目录下，放置增强执行代码，即
augment.py

import os
import cv2
from imgaug import augmenters as iaa

# 定义数据增强的方法，这里使用了 imgaug 库
seq = iaa.Sequential([
    iaa.Fliplr(0.5),  # 左右翻转
    iaa.Affine(rotate=(-10, 10))  # 旋转
])

input_dir = 'datasets'
output_dir = 'aug_datasets'

if not os.path.exists(output_dir):
    os.makedirs(output_dir)

for root, dirs, files in os.walk(input_dir):
    for subdir in dirs:
        subdir_path = os.path.join(root, subdir)
        output_subdir = subdir_path.replace(input_dir, output_dir)
        if not os.path.exists(output_subdir):
            os.makedirs(output_subdir)

        for file in os.listdir(subdir_path):
            if file.lower().endswith(('.png', '.jpg', '.jpeg', 'JPG')):
                image_path = os.path.join(subdir_path, file)
                image = cv2.imread(image_path)

                # 数据增强
                augmented_images = [seq.augment_image(image) for _ in range(3)]  # 扩充到原来的3倍，包括原图像

                for idx, img in enumerate(augmented_images):
                    output_path = os.path.join(output_subdir, f"{file.split('.')[0]}_aug{idx}.jpg")
                    cv2.imwrite(output_path, img)

执行完毕后，新的文件夹aug_datasets里，就有着增强之后的图像