DataAugment:同时对image和mask进行变换

博主想使用Unet网络完成一个分割任务,手边只有40张图和对应的mask,需要进行data augment.
做数据增强有很多工具,常用的是使用keras内置的ImageDataGenerator生成器生成图片,但是这个工具只能对一张图进行随机变化,而image和mask是一一对应的,二者必须同时进行同种变化.

下面隆重介绍一个强大的数据增强工具augmentor

它的使用方法十分简单

安装Augmentor

pip install Augmentor

对图片进行随机旋转

import Augmentor
p = Augmentor.Pipeline("/path/to/images")
p.rotate(probability=1, max_left_rotation=5, max_right_rotation=5) #probability表示以一定概率随机处理图片
p.sample(500) #产生500张图片

image and ground truth data can be identically augmented

p = Augmentor.Pipeline("/path/to/images")
# Point to a directory containing ground truth data.
# Images with the same file names will be added as ground truth data
# and augmented in parallel to the original data.
p.ground_truth("/path/to/ground_truth_images")
# Add operations to the pipeline as normal:
p.rotate(probability=1, max_left_rotation=5, max_right_rotation=5)
p.flip_left_right(probability=0.5)
p.zoom_random(probability=0.5, percentage_area=0.8)
p.flip_top_bottom(probability=0.5)
p.sample(50)

在旋转图片时,常常会在图片周围产生空白填充,如图

遇到这种情况,Augmentor会在旋转的时候同时缩放图片,不致在四周出现黑色填充

问题

在使用
ground_truth()函数时,如果路径中有多张图片,将会导致augment之后的mask和image不对应,因此只能在路径中存放一张图片,如果有很多组数据需要augment则需要将他们单个存放在文件夹中

下面是我的代码:

 # -*- coding: utf-8 -*-
import Augmentor
import glob
import os
import random
from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img

train_path = 'train'
groud_truth_path = 'mask'
img_type = 'jpg'
train_tmp_path = 'tmp/train'
mask_tmp_path = 'tmp/mask'

def start(train_path,groud_truth_path):
    train_img = glob.glob(train_path+'/*.'+img_type)
    masks = glob.glob(groud_truth_path+'/*.'+img_type)

    if len(train_img) != len(masks):
        print ("trains can't match masks")
        return 0
    for i in range(len(train_img)):
        train_img_tmp_path = train_tmp_path + '/'+str(i)
        if not os.path.lexists(train_img_tmp_path):
            os.mkdir(train_img_tmp_path)
        img = load_img(train_path+'/'+str(i)+'.'+img_type)
        x_t = img_to_array(img)
        img_tmp = array_to_img(x_t)
        img_tmp.save(train_img_tmp_path+'/'+str(i)+'.'+img_type)

        mask_img_tmp_path =mask_tmp_path +'/'+str(i)
        if not os.path.lexists(mask_img_tmp_path):
            os.mkdir(mask_img_tmp_path)
        mask = load_img(groud_truth_path+'/'+str(i)+'.'+img_type)
        x_l = img_to_array(mask)
        mask_tmp = array_to_img(x_l)
        mask_tmp.save(mask_img_tmp_path+'/'+str(i)+'.'+img_type)
        print ("%s folder has been created!"%str(i))
    return i+1


def doAugment(num):
    sum = 0
    for i in range(num):
        p = Augmentor.Pipeline(train_tmp_path+'/'+str(i))
        p.ground_truth(mask_tmp_path+'/'+str(i))
        p.rotate(probability=0.5, max_left_rotation=5, max_right_rotation=5)#旋转
        p.flip_left_right(probability=0.5)#按概率左右翻转
        p.zoom_random(probability=0.6, percentage_area=0.99)#随即将一定比例面积的图形放大至全图
        p.flip_top_bottom(probability=0.6)#按概率随即上下翻转
        p.random_distortion(probability=0.8,grid_width=10,grid_height=10, magnitude=20)#小块变形
        count = random.randint(40, 60)
        print("\nNo.%s data is being augmented and %s data will be created"%(i,count))
        sum = sum + count
        p.sample(count)
        print("Done")
    print("%s pairs of data has been created totally"%sum)


a = start(train_path, groud_truth_path)
doAugment(a)