imgaug增强图像专门针对语义分割

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):
        img_prefix = image_name.split("\\")[-1].split(".jpg")[0]
        # 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 + img_prefix + name + ".jpg")
        segmap.save(self.SegmentationClass_aug_dir + img_prefix + 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 + img_prefix + 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 + img_prefix + name + ".png")
            self.image_num += 1


if __name__ == "__main__":
    import os
    root_img_dir = r"H:\image"
    root_mask_dir = r"H:\mask"
    root_save_img_dir = r"H:\img_aug" + "\\"
    root_save_mask_dir = r"H:\mask_aug" + "\\"

    for img_name in os.listdir(root_img_dir):
        image_name = os.path.join(root_img_dir, img_name)
        segmap_name = os.path.join(root_mask_dir, img_name.replace(".jpg", ".png"))

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