数据增强之王——Augmentor

其实我们都知道Pytorch的torchvision的transforms并非真正意义上的数据增强，它只是将图像做了变换，实际上图像数据集的总量是不变的。

现在我们使用Augmentor来做真正意义上的数据增强，首先安装Augmentor

pip install Augmentor

如果我们的ground truth，即mask图像是单通道的，需要转成3通道的RGB图像

from PIL import Image
import os

if __name__ == '__main__':

    img_path = "/home/user/line/data_dataset_voc/SegmentationClassPNG/"
    for im_name in os.listdir(img_path):
        img_mode = Image.open(img_path + im_name)
        print(img_mode.mode)
        img = img_mode.convert('RGB')
        img.save(img_path + im_name)

由于一般mask图像是png格式的，Augmentor需要保证原始图像和mask图像格式一致，我这里原始图像是jpg格式的

import cv2
import os

if __name__ == '__main__':

    img_path = "/home/user/line/data_dataset_voc/JPEGImages/"
    output_path = "/home/user/line/data_dataset_voc/JPEGImagespng/"
    for im_name in os.listdir(img_path):
        print(img_path + im_name)
        img = cv2.imread(img_path + im_name)
        im_name_new = im_name.split('.')[0]+'.png'
        cv2.imwrite(output_path + im_name_new, img)

最后就是开始做数据增强处理了，当然我这里只做了旋转和翻转处理，你还可以作一些其他的数据增强。最后会得到20000张增强图像。

import Augmentor

if __name__ == '__main__':

    p = Augmentor.Pipeline("/home/user/line/data_dataset_voc/JPEGImagespng/")
    p.ground_truth("/home/user/line/data_dataset_voc/SegmentationClassPNG/")
    p.rotate(probability=0.8, max_left_rotation=25, max_right_rotation=25)
    p.flip_left_right(probability=0.5)
    p.flip_top_bottom(probability=0.5)
    p.sample(20000)

最后再分别拷贝到原始图像数据集以及mask数据集中即可开始训练。

{{o.name}}

{{m.name}}