はじめに

以前にImageDataGeneratorを使ってData Augmentation（水増し）を行ったのですが、ImageDataGeneratorが持っていない水増し方法も使いたいと思っていました。今回それを実現してみました。

環境

Google Colaboratory
TensorFlow 2.0 Alpha

コード

こちらです。

コード解説

import numpy as np
from tensorflow.keras.preprocessing.image import ImageDataGenerator

class CustomImageDataGenerator(ImageDataGenerator):
    def __init__(self, cutout_mask_size = 0, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.cutout_mask_size = cutout_mask_size

    def cutout(self, x, y):
        return np.array(list(map(self._cutout, x))), y

    def _cutout(self, image_origin):
        # 最後に使うfill()は元の画像を書き換えるので、コピーしておく
        image = np.copy(image_origin)
        mask_value = image.mean()

        h, w, _ = image.shape
        # マスクをかける場所のtop, leftをランダムに決める
        # はみ出すことを許すので、0以上ではなく負の値もとる(最大mask_size // 2はみ出す)
        top = np.random.randint(0 - self.cutout_mask_size // 2, h - self.cutout_mask_size)
        left = np.random.randint(0 - self.cutout_mask_size // 2, w - self.cutout_mask_size)
        bottom = top + self.cutout_mask_size
        right = left + self.cutout_mask_size

        # はみ出した場合の処理
        if top < 0:
            top = 0
        if left < 0:
            left = 0

        # マスク部分の画素値を平均値で埋める
        image[top:bottom, left:right, :].fill(mask_value)
        return image

    def flow(self, *args, **kwargs):
        batches = super().flow(*args, **kwargs)

        # 拡張処理
        while True:
            batch_x, batch_y = next(batches)

            if self.cutout_mask_size > 0:
                result = self.cutout(batch_x, batch_y)
                batch_x, batch_y = result                        

            yield (batch_x, batch_y)     

datagen = CustomImageDataGenerator(rotation_range=10, horizontal_flip=True, zoom_range=0.1, cutout_mask_size=16)

ImageDataGeneratorを継承したClassを作成し、flowメソッドをオーバーライドし、そこからcutoutを呼び出しています。

出力結果

ちゃんとcutoutが入っているようです。

参考にしたページ