imgaug成对图像数据增强（附代码）

本文使用imgaug库集成了一些主要的成对图像增强方法，如图像翻转、旋转、平移、剪切等（共13中）。

1、增强效果部分

翻转

旋转

平移

加噪声

剪切

代码

from imgaug import augmenters as iaa #引入数据增强的包
import os
import cv2
from matplotlib import pyplot as plt
import numpy as np
from imgaug.augmentables.segmaps import SegmentationMapsOnImage

def FlipHorizontalImage():
    """水平翻转"""
    sequential.append(iaa.Fliplr())

def FlipVerticalImage():
    """垂直翻转"""
    sequential.append(iaa.Flipud())

def AffineRotateImage(r=(-180,180)):
    """旋转：默认随机旋转(-180,180)"""
    sequential.append(iaa.Affine(rotate=r))

def AffineTranslationImage(t=(-0.5,0.5)):
    """平移,默认平移(-50%,50%)"""
    sequential.append(iaa.Affine(translate_percent=t))

def AffineScaleImage(s=(0.4,0.8)):
    """缩放,默认缩放(0.4倍,0.8倍)"""
    sequential.append(iaa.Affine(scale=s))

def CropImage(p=(0, 0.5)):
    """剪切，默认剪切幅度(0,50%)"""
    sequential.append(iaa.Crop(percent=p))

def AddGaussianNoiseImage(s=(0, 0.2*255)):
    """加入高斯噪声"""
    sequential.append(iaa.AdditiveGaussianNoise(scale=s))

def FindMaxFile(results_file_list):
    """找到最大的文件"""
    flag = 1
    for file_name in results_file_list:
        if int(file_name[3:])>flag:
            flag = int(file_name[3:])
    return 'exp'+str(flag)

def PiecewiseAffineImage(s=(0.05,0.15)):
    """对图像进行扭曲"""
    sequential.append(iaa.PiecewiseAffine(scale=s))

def AffineShearImage(s=(-50, 50)):
    """进行剪切"""
    sequential.append(iaa.Affine(shear=s))

def AddRungImage(i=4):
    """加横杠,每i条线加一个横杠"""
    def img_func(images, random_state, parents, hooks):
        for img in images:
            img[::i] = 0
        return images
    def keypoint_func(keypoints_on_images, random_state, parents, hooks):
        return keypoints_on_images
    
    sequential.append(iaa.Lambda(img_func, keypoint_func))

def AddGaussianBlurImage(sigma=5.0):
    """高斯模糊"""
    sequential.append(iaa.GaussianBlur(sigma=(0.0, sigma)))

def DropoutImage(p=0.2):
    """丢弃一些像素,p为丢弃比例"""
    sequential.append(iaa.Dropout(p=(0, p)))

def AddAverageBlurImage(k=(5,5)):
    """平均模糊"""
    sequential.append(iaa.AverageBlur(k=k))

if __name__ == '__main__':
    pwd = os.getcwd()
    file_path = os.path.join(pwd,'datasets','train') # 原图片存储路径
    file_name_list = os.listdir(file_path) # 该路径下所有的文件名字
    label_path = os.path.join(pwd,'datasets','label') # 标签数据

    # 创建数据增强保存的文件夹
    results_path = os.path.join(pwd,'results')
    results_file_list = os.listdir(results_path)
    if len(results_file_list) == 0:
        results_path = os.path.join(results_path, 'exp0')
        os.mkdir(os.path.join(results_path))
        os.mkdir(os.path.join(results_path,'train'))
        os.mkdir(os.path.join(results_path,'label'))
        os.mkdir(os.path.join(results_path, 'show'))
    else:
        file_name = FindMaxFile(results_file_list)
        num = int(file_name[3:])
        num += 1
        file_name = 'exp'+str(num)
        results_path = os.path.join(results_path, file_name)
        os.mkdir(results_path)
        os.mkdir(os.path.join(results_path, 'train'))
        os.mkdir(os.path.join(results_path, 'show'))
        os.mkdir(os.path.join(results_path, 'label'))


    image_result_list = [] # 放结果图片数据的集合

    sequential = iaa.Sequential([
    ], random_order=True)

    """翻转"""
    FlipHorizontalImage() # 水平翻转
    # FlipVerticalImage() # 垂直翻转
    """放射变换"""
    # AffineRotateImage() # 旋转
    # AffineTranslationImage() # 平移
    # AffineScaleImage() # 缩小
    # PiecewiseAffineImage() # 扭曲
    # AffineShearImage() # 剪切
    """裁剪"""
    # CropImage() # 裁剪
    """加噪声"""
    # AddGaussianNoiseImage() # 加入高斯噪声
    # AddRungImage() # 加横杠
    # DropoutImage() # 丢弃
    """模糊"""
    # AddGaussianBlurImage() # 加入高斯模糊
    # AddAverageBlurImage() # 平均模糊

    for file_name in file_name_list:
        img = cv2.imread(os.path.join(file_path, file_name))
        img_label = cv2.imread(os.path.join(label_path, file_name))
        img_label = SegmentationMapsOnImage(img_label, shape=img.shape)

        image_aug, label_aug = sequential(image=img, segmentation_maps=img_label)
        cv2.imwrite(os.path.join(results_path,'train',file_name), image_aug)
        cv2.imwrite(os.path.join(results_path, 'label', file_name), label_aug.get_arr_int())

        # BGR->RGB
        img_vstack = np.hstack((
                                           np.vstack((img[:,:,::-1], image_aug[:,:,::-1])),
                                           np.vstack((img_label.get_arr_int()[:, :, ::-1],
                                                      label_aug.get_arr_int()[:, :, ::-1]))
                                           ))

        cv2.imwrite(os.path.join(results_path, 'show', file_name), img_vstack[:,:,::-1])
        image_result_list.append(img_vstack)

    for i,img in enumerate(image_result_list):
        plt.imshow(img)
        plt.show()

使用方法

目录结构

找到main.py中的这几行，你要使用哪些增强方法就不要注释那几行，例如要使用水平翻转+旋转+平移，就不要注释这几行

 """翻转"""
    FlipHorizontalImage() # 水平翻转
    # FlipVerticalImage() # 垂直翻转
    """放射变换"""
    AffineRotateImage() # 旋转
    AffineTranslationImage() # 平移
    # AffineScaleImage() # 缩小
    # PiecewiseAffineImage() # 扭曲
    # AffineShearImage() # 剪切
    """裁剪"""
    # CropImage() # 裁剪
    """加噪声"""
    # AddGaussianNoiseImage() # 加入高斯噪声
    # AddRungImage() # 加横杠
    # DropoutImage() # 丢弃
    """模糊"""
    # AddGaussianBlurImage() # 加入高斯模糊
    # AddAverageBlurImage() # 平均模糊

运行main.py，在results文件夹中看结果