有一個需求

使用imgaug工具包進行影像分割資料增強


說明：輸入有兩張圖片，分別是RGB圖片和對應的label圖片，然後通過程式碼對這兩張圖片做了 縮放、映象+上下翻轉、旋轉、xy平移、裁剪、旋轉 + 裁剪、高斯平滑的影像增強


兩張輸入圖片：
RGB圖片

對應的label

程式碼如下

import random
import glob
import numpy as np
import imgaug as ia
import imgaug.augmenters as iaa
from imgaug.augmentables.segmaps import SegmentationMapsOnImage
from PIL import Image


class ImageAugmentor(object):
    def __init__(self, image_dir=None, segmap_dir=None, image_aug_dir=None, SegmentationClass_aug_dir=None):
        self.image_dir = image_dir  # 存放原圖的目錄，必須是RGB
        self.segmap_dir = segmap_dir  # 存放原圖對應的標籤，必須是p模式的圖片
        self.image_aug_dir = image_aug_dir
        self.SegmentationClass_aug_dir = SegmentationClass_aug_dir

        self.image_num = 1
        self.seed_set()

    def seed_set(self, seed=1):
        np.random.seed(seed)
        random.seed(seed)
        ia.seed(seed)

    def array2p_mode(self, alpha_channel):
        """alpha_channel is a binary image."""
        assert set(alpha_channel.flatten().tolist()) == {0, 1}, "alpha_channel is a binary image."
        alpha_channel[alpha_channel == 1] = 128
        h, w = alpha_channel.shape
        image_arr = np.zeros((h, w, 3))
        image_arr[:, :, 0] = alpha_channel
        img = Image.fromarray(np.uint8(image_arr))
        img_p = img.convert("P")
        return img_p

    def augmentor(self, image):
        height, width, _ = image.shape
        resize = iaa.Sequential([
            iaa.Resize({"height": int(height/2), "width": int(width/2)}),
        ])

        fliplr_flipud = iaa.Sequential([
            iaa.Fliplr(),
            iaa.Flipud(),
        ])

        guassian_blur = iaa.Sequential([
            iaa.GaussianBlur(sigma=(1.5, 2.5)),
        ])

        rotate = iaa.Sequential([
            iaa.Affine(rotate=(-90, 90))
        ])

        translate = iaa.Sequential([
            iaa.Affine(translate_percent=(0.2, 0.5))
        ])

        crop_and_pad = iaa.Sequential([
            iaa.CropAndPad(percent=(-0.25, 0), keep_size=False),
        ])

        rotate_and_crop = iaa.Sequential([
            iaa.Affine(rotate=45),
            iaa.CropAndPad(percent=(-0.25, 0), keep_size=False)
        ])

        ops = [resize, fliplr_flipud, rotate, translate, crop_and_pad, rotate_and_crop, guassian_blur]
        #    縮放、映象+上下翻轉、旋轉、xy平移、裁剪、旋轉 + 裁剪、高斯平滑
        return ops

    def augment_image(self, image_name, segmap_name):
        # 1.Load an image.
        image = Image.open(image_name)
        segmap = Image.open(segmap_name)

        name = f"{self.image_num:04d}"
        image.save(self.image_aug_dir + name + ".jpg")
        segmap.save(self.SegmentationClass_aug_dir + name + ".png")
        self.image_num += 1

        image = np.array(image)
        segmap = SegmentationMapsOnImage(np.array(segmap), shape=image.shape)

        # 2. define the ops
        ops = self.augmentor(image)

        # 3.execute ths ops
        for _, op in enumerate(ops):
            name = f"{self.image_num:04d}"
            print(f"當前增強了{self.image_num:04d}張資料...")
            images_aug_i, segmaps_aug_i = op(image=image, segmentation_maps=segmap)
            images_aug_i = Image.fromarray(images_aug_i)
            images_aug_i.save(self.image_aug_dir + name + ".jpg")

            segmaps_aug_i_ = segmaps_aug_i.get_arr()
            segmaps_aug_i_ = self.array2p_mode(segmaps_aug_i_)
            segmaps_aug_i_.save(self.SegmentationClass_aug_dir + name + ".png")
            self.image_num += 1


if __name__ == "__main__":
    image_name = "./JPEGImages/0036.jpg"
    segmap_name = "./SegmentationClass/0036.png"
    image_aug_dir = "./JPEG_AUG/"
    SegmentationClass_aug_dir = "./SegmentationClass_AUG/"

    image_augmentation = ImageAugmentor(image_aug_dir=image_aug_dir, SegmentationClass_aug_dir=SegmentationClass_aug_dir)
    image_augmentation.augment_image(image_name, segmap_name)

實現方法

藉助imgaug工具包實現


引數說明
image_name
segmap_name
image_aug_dir
SegmentationClass_aug_dir

結果展示

原圖增強後的七張圖

標籤增強後的七張圖