深度学习常用图像数据增强库albumentations系列教程（二）

用户9875047

发布于 2022-12-07 10:21:00

7790

发布于 2022-12-07 10:21:00

文章被收录于专栏：机器视觉全栈er机器视觉全栈er

目标：掌握更加复杂和自动化的图像数据增强方法

在第一篇教程中，我们讲述了简单的图像增强技巧，本节通过使用Compose方法将各种图像增强的方法组合起来，形成一个图像增强的pipeline，方便产生大量不同种类的增强图片。

首先我们来看下使用compose方法对图片进行增强的示例：

import random
import cv2
from matplotlib import pyplot as plt
import albumentations as A

def visualize(image):
    plt.figure(figsize=(10, 10))
    plt.axis('off')
    plt.imshow(image)
    plt.show()

image = cv2.imread('assets/cvtutorials.png')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

transform = A.Compose([
    A.CLAHE(),
    A.RandomRotate90(),
    A.Transpose(),
    A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.50, rotate_limit=45, p=.75),
    A.Blur(blur_limit=3),
    A.OpticalDistortion(),
    A.GridDistortion(),
])
random.seed(100)
augmented_image = transform(image=image)['image']
visualize(augmented_image)

上述代码会对原始图像做出7种一定概率的随机变换，然后生成一种图像，如：

现在，我们去掉增强图像可视化代码，并生成大量增强图片,图片会保存在res文件夹下，代码如下：

import random
import cv2
from matplotlib import pyplot as plt
import albumentations as A

GENERATED_PICS_SIZE = 2 # 可以设置想生成的图片个数
RES_PATH_PREFIX = "res"

image = cv2.imread('assets/cvtutorials.png')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

for i in range(GENERATED_PICS_SIZE):
    res_path = RES_PATH_PREFIX + "\\" + str(i) + ".jpg"
    transform = A.Compose([
        A.CLAHE(),
        A.RandomRotate90(),
        A.Transpose(),
        A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.50, rotate_limit=45, p=.75),
        A.Blur(blur_limit=3),
        A.OpticalDistortion(),
        A.GridDistortion(),
    ])
    augmented_image = transform(image=image)['image']
    image = cv2.cvtColor(augmented_image, cv2.COLOR_RGB2BGR)
    cv2.imencode('.jpg', image)[1].tofile(res_path)

文件同一目录下记得创建文件夹res，否则会报错，也可以自己添加代码，检测不到文件夹路径就创建文件夹，自己思考 :p

仅仅这样还是不够的，我们常常需要对某个文件夹下所有的图片进行增强，所以又要加入遍历文件夹下所有图片的功能，代码如下：

import random
import cv2
import albumentations as A
import glob
import os

GENERATED_PICS_SIZE = 2
RES_PATH_PREFIX = "res"
PICS_PATH = 'assets'

paths = glob.glob(os.path.join(PICS_PATH, '*.jpg'))
for i in range(len(paths)):
    pic_path = paths[i]
    pic_path_prefix = pic_path.split('.')[0].split('\\')[-1]

    image = cv2.imread(pic_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    for i in range(GENERATED_PICS_SIZE):
        res_path = RES_PATH_PREFIX + "\\" + pic_path_prefix + str(i) + ".jpg"
        transform = A.Compose([
            A.CLAHE(),
            A.RandomRotate90(),
            A.Transpose(),
            A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.50, rotate_limit=45, p=.75),
            A.Blur(blur_limit=3),
            A.OpticalDistortion(),
            A.GridDistortion(),
        ])
        augmented_image = transform(image=image)['image']
        image = cv2.cvtColor(augmented_image, cv2.COLOR_RGB2BGR)
        cv2.imencode('.jpg', image)[1].tofile(res_path)

代码文件所在目录有文件夹assets和res，里面分别是原始图像和生成后图像，这样我们实现了批量自动生成大量图片。