Albumentationの拡張方法 - のんびりしているエンジニアの日記

皆さんこんにちは
お元気ですか。週末はALASKA2の準備をはじめてみました。

さて、今回のAlbumentationの拡張方法を書いておきます。
AlbumentationはKaggleなどの画像コンペでよく利用されるData Augmentation（データ拡張）のライブラリです。
Albumentationは様々なシチュエーションを想定して継承元のクラスがいくつかあります。
これらを継承し実装することで新しいAugmentationのクラスを開発できます。

今回はどうすればAlbumentationで新しいData Augmentationのクラスを作成できるのかを記載します。

Albumentationの構造

Albumentationは次の2つの基底となるクラスが存在します。

ImageOnlyTransform・・画像のみに処理を適用する
DualTransform・・画像＋物体検出（＋その他）に処理を適用する

新しいData Augmentationのクラスを作成する場合これらを継承して作ります。
具体的に作成する必要があるのは次のメソッドです。
ただし、コンペで利用する場合、全てを実装する必要はなく、この中から必要なものを実装すれば良いでしょう。

メソッド	説明
__init__	パラメータ
get_params	パラメータを取得する。ランダムで引数を取得するのはこの箇所
apply	画像に対して適用する処理を実装する
apply_to_mask	マスクに対して適用する処理を実装する（Segmentation系限定）
apply_to_masks	複数のマスクに対して適用する処理を実装する（Segmentation系限定）
apply_to_bbox	バウンディングボックスに対して適用する処理を実装する（物体検出系限定）
apply_to_keypoints	キーポイントに対して適用する処理を実装する

Albumentstionの新しいクラスを作る

では、Albumentationのクラスを作成してみましょう。
今回はRandomErasingと呼ばれる画像の一部分を置き換える（削除する）拡張です。
applyのみ実装すれば可能ですので、最初のサンプルとしては便利かと思います。

実装は次の通りです。

class RandomErasing(ImageOnlyTransform):
    def __init__(self, always_apply=False, p=0.5, sl=0.02, sh=0.4, r1=0.3, mean=[0.4914, 0.4822, 0.4465]):
        super(RandomErasing, self).__init__(always_apply=always_apply, p=p) # (1) 継承元のクラスへ渡すパラメータ（＝おまじまい）
        self.mean = mean # (2)パラメータ
        self.sl = sl
        self.sh = sh
        self.r1 = r1

    def apply(self, img, **params):
        # (3) 実行
        for attempt in range(100):
            area = img.shape[0] * img.shape[1]

            target_area = random.uniform(self.sl, self.sh) * area
            aspect_ratio = random.uniform(self.r1, 1 / self.r1) 

            h = int(round(math.sqrt(target_area * aspect_ratio)))
            w = int(round(math.sqrt(target_area / aspect_ratio)))

            if w < img.shape[1] and h < img.shape[0]:
                x1 = random.randint(0, img.shape[0] - h)
                y1 = random.randint(0, img.shape[1] - w)
                if len(img.shape) == 3:
                    img[x1:x1 + h, y1:y1 + w, 0] = self.mean[0]
                    img[x1:x1 + h, y1:y1 + w, 1] = self.mean[1]
                    img[x1:x1 + h, y1:y1 + w, 2] = self.mean[2]
                else:
                    img[x1:x1 + h, y1:y1 + w] = self.mean[0]
                return img

        return img

(1) 継承元のクラスへパラメータを渡す事が必要です。

always_applyはこのパラメータがTrueの場合に常に適用することを示し、pはこのData Augmentationを実行する確率を示します。

(2) パラメータの設定

コンストラクタ側にパラメータを設定しています。
これらの変数はapplyメソッドで参照します。

(3) 実行

RandomErasingの実装部分になります。
ここからはRandomErasingの実装になるので詳細は割愛しますが、通常の実装をこなせば問題ありません。
このメソッドの返り値は処理後の画像です。

後は通常通りAlbumentationに組み込んであげれば、Albumentationの世界で動作します。

最後に

私も頻繁にこのライブラリを利用しますが、やはり、Albumentation上で処理できるとやはり便利です。
新しいアルゴリズムを簡単に試すことも可能なのでぜひ、試してみてください。