データ拡張ライブラリAugmentorの画像変換機能だけを利用する

データ拡張ライブラリとして有名なAugmentor、色々な変換がサポートされていて良いのですが、何故かデータ拡張した画像をファイル出力させたがるフシがあります。KerasやPyTorch向けのGeneratorもあるのですが、ここのIssueにあるように、(入力画像, ラベル) のペアしか対応しておらず、(入力画像, 出力画像) のようにsegmentationのマスク画像のようなものを（公式のGeneratorでは）扱うことができません。

画像変換機能だけ提供してくれてれば良いのですが、丁度良さげな関数がないようです。一応 p.torch_transform() で1枚画像を変換する関数を取ってこれる（PyTorch専用というわけではない）のですが、(入力画像, マスク画像) といったペア画像を同じ変換で変換することができません。

ということで、p.torch_transform() を参考に画像変換を行う関数を作ってみます。Augmentor内部では画像はPIL.Imageである前提があるのでNumPy配列から変換していますが、最初からPIL.Imageを使う場合は変換は不要です。あと、本来はBGR->RGB変換もしないといけないのですが、色に関わる変換を使わなければ問題ないので、そのままにしています。ちなみに色（やコントラスト）に関する変換もマージ済みのこのPRで実装されています。

import random
from PIL import Image
import numpy as np
import cv2
import Augmentor


def get_transform_func():
    # 変換のパイプラインはここで定義する
    p = Augmentor.Pipeline()
    p.flip_left_right(probability=0.5)
    p.rotate(probability=1, max_left_rotation=5, max_right_rotation=5)
    p.zoom_random(probability=1, percentage_area=0.95)

    def transform_images(images):
        images = [Image.fromarray(image) for image in images]
        for operation in p.operations:
            r = round(random.uniform(0, 1), 1)
            if r <= operation.probability:
                images = operation.perform_operation(images)
        images = [np.array(image) for image in images]
        return images
    return transform_images


def main():
    target = np.zeros((256, 512, 3), dtype=np.uint8)
    transform_images = get_transform_func()

    while True:
        image1 = np.zeros((256, 256, 3), dtype=np.uint8)
        cv2.rectangle(image1, (100, 100), (200, 200), (255, 0, 0), -1)
        image2 = np.zeros((256, 256, 3), dtype=np.uint8)
        cv2.rectangle(image2, (100, 100), (200, 200), (255, 255, 255), -1)
        images = [image1, image2]
        images = transform_images(images)
        target[:, :256] = images[0]
        target[:, 256:] = images[1]

        cv2.imshow("transformed images", target)
        key = cv2.waitKey(-1)

        # "q": quit
        if key == 113:
            return 0


if __name__ == '__main__':
    main()

え？imgaug？albumentations？