1. 程式人生 > >資料增廣 | 原圖 和mask 同時增加——適用語義分割等

資料增廣 | 原圖 和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)