python图像数据增强——imgaug (三)

(一)imgaug基础用法
(二)imgaug进阶示例
(三)imgaug图像分割数据增强

之前两章展示了最基本的imgaug的基础用法,以及一个较为全面的各种增强方法的展示,这里我们介绍在分割问题中的数据增强方法。

在分类问题中,输入为一张图像,输出为一个类别,在这种问题中,我们只需要对图像进行增强,并不需要对标签做出改动。而在分割问题中,当训练使用的图像被翻转或者扭曲之后,标签也要对应的做出相同的翻转或者扭曲,因此在增强的过程中不能仅仅对图像进行处理,同时也要对分割的标签进行相同的处理。

为了对图像和标签做出相同的处理,我们需要确定对一张图的增强方法,然后先后应用在图像和标签上面,这里我们就需要使用该函数seq.to_deterministic()得到一个确定的增强函数。要想知道如何使用,首先我们就要有两个分别记录原图和分割标签的变量,分别叫做image,和segmap。

实现的方法如下,我们这里使用系统自带的一张图片进行增强,这里是在原网站教程的基础上进行改进的,因为在手册中并没有一个中间步骤的输出可以用来训练网络,缺少一步把 SegmentationMapOnImage类型的数据转换成numpy类型数据的步骤,这里我们增加了这一步。我们最后得到的结果就如下图所示。最后再补充一句就是这里对分割增强的方法在这个imgaug库中还在开发中,有些测试还没有做完全,不过目前可以正常使用的。
python图像数据增强——imgaug (三)

import imgaug as ia
from imgaug import augmenters as iaa
import imageio
import numpy as np

ia.seed(1)
image = ia.quokka(size=(128, 128), extract="square") #加载示例图像进来,大小为(128,128,3)

#创建一个分割的图
segmap = np.zeros((128, 128), dtype=np.int32)
segmap[28:71, 35:85] = 1
segmap[10:25, 30:45] = 2
segmap[10:25, 70:85] = 3
segmap[10:110, 5:10] = 4
segmap[118:123, 10:110] = 5

#将图片转换为SegmentationMapOnImage类型
segmap = ia.SegmentationMapOnImage(segmap, shape=image.shape, nb_classes=1+5)


# 定义数据增强方法
seq = iaa.Sequential([
    iaa.Dropout([0.05, 0.2]),      # drop 5% or 20% of all pixels
    iaa.Sharpen((0.0, 1.0)),       # sharpen the image
    iaa.Affine(rotate=(-45, 45)),  # rotate by -45 to 45 degrees (affects heatmaps)
    iaa.ElasticTransformation(alpha=50, sigma=5)  # apply water effect (affects heatmaps)
], random_order=True)

images_aug = []
segmaps_aug = []

#这里可以通过加入循环的方式,对多张图进行数据增强。
seq_det = seq.to_deterministic() #确定一个数据增强的序列
images_aug = seq_det.augment_image(image) #将方法应用在原图像上
segmaps_aug = seq_det.augment_segmentation_maps([segmap])[0].get_arr_int().astype(np.uint8)
# 将方法应用在分割标签上,并且转换成np类型

segmaps_aug = ia.SegmentationMapOnImage(segmaps_aug, shape=image.shape, nb_classes=1+5)
# 将分割结果转换为SegmentationMapOnImage类型,方便后面可视化

#可视化
cells = []
cells.append(image)
cells.append(segmap.draw_on_image(image))
cells.append(images_aug)
cells.append(segmaps_aug.draw_on_image(images_aug))
cells.append(segmaps_aug.draw(size=images_aug.shape[:2]))  
grid_image = ia.draw_grid(cells, cols=5)
imageio.imwrite("example_segmaps.jpg", grid_image)