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
里,就有着增强之后的图像