セグメンテーションで反転、回転の水増しを実装する方法

こんにちは、タナカです。

この記事では、セマンティックセグメンテーションで反転、回転の水増しを実施する方法について紹介します。

pytorchにtransformという水増しを実施できるクラスがあるものの、学習データと教師データがうまく紐づけできなかったので、他の方法を試しました。

試した方法として、albumentationsと呼ばれるライブラリを使うことで解決できたのでその内容を記述したいと思います。

この記事の内容

セマンティックセグメンテーションで反転、回転の水増しを実施する方法

1. セマンティックセグメンテーションとは
2. 画像の水増しとは
3. 画像の水増しをするためのライブラリ
まとめ

1. セマンティックセグメンテーションとは

セマンティックセグメンテーションとは、ピクセル単位で各物体を検出する機械学習アルゴリズムの一つです。

自動運転、医療用画像処理、工場での製品検査などに応用されています。

セマンティックセグメンテーションで学習を行う場合、下図のような画像を扱うことになります。左が学習データ、右が教師データです。

このような画像の組み合わせを何枚も準備して、UNetやSegNetと呼ばれるモデルを使うことで学習を行うことができます。

2. 画像の水増しとは

画像の水増しとは、元々持っている画像データに何らかの画像処理を加えることで別の画像データを生成することを言います。

別のデータを生成することで元々10枚しかもっていなかったとしても、画像データを2倍、5倍と数を増やすことができます。

具体的な例でいうと、同じ画像であっても10°回転させた画像は違う画像として扱うことができるということです。

3. 画像の水増しをするためのライブラリ

私が普段使っているライブラリは2つあります。

transforms
albumentations

transformsについては、【pytorch】transformを実装する方法で書いていますので参考にしてください。

albumentationsは、外部ライブラリなのでpip等でインストールする必要があります。

pip install albumentations

参考: 公式リファレンス

これらのライブラリを使って、セマンティックセグメンテーションの水増しを実装してみました。

結果としては、transformsでは反転や回転などの処理をうまく実行することができませんでした。一方、albumentationsでは正しく処理することができました。

初めに、transformsでの実行結果がこちらになります。

実行したコードがこちらになります。

from torchvision import transforms
import matplotlib.pyplot as plt
from PIL import Image

# 画像を表示する関数
def imgShow(imgList, row_num = 1):
    img_num = len(imgList)
    if img_num == 1:
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.imshow(imgList[0])
    else:
        fig, axes = plt.subplots(row_num, round(img_num / row_num))
        ax = axes.ravel()
        for i in range(img_num):
            ax[i].imshow(imgList[i])
            ax[i].xaxis.set_visible(False)
            ax[i].yaxis.set_visible(False)
    plt.show()

# transformによる水増し関数
def trans_aug():
    transform = transforms.Compose([
        transforms.RandomHorizontalFlip(p = 0.5),
    ])
    return transform

# 画像の読み込み
img_pil = Image.open('apple-1834639_640.jpg')
mask_pil = Image.open('apple-1834639_640_mask.jpg')

# 水増し処理
transform = trans_aug()
img = transform(img_pil)
mask = transform(mask_pil)

# 画像表示
imgShow([img, mask])

学習データと教師データが対になっているので、学習データが水平方向に反転するなら、教師データも水平方向に反転する必要があります。

しかし、transformsではそれぞれの画像を読み込ませる必要があるため、教師データは反転するけど、学習データは反転しないといったことが起きます。ただ、もしかしたらtransformsでも実施できる方法があるかもしれません。

transformsの注意点として、読み込ませる画像はPIL形式にする必要があります。

また、transforms.RandomHorizontalFlip(p = 0.5)のpは、probabilityを示しており50%の確率で水平方向に反転することを意味しています。

次に、albumentationsでの実行結果がこちらになります。期待通りの結果にしかならないので実行結果のみ載せてます。

実行したコードがこちらになります。

import cv2
import albumentations as albu

# albumentationsによる水増し関数
def albu_aug():
    augmentation =albu.Compose([
        albu.HorizontalFlip(p = 0.5),
    ])
    return augmentation

# 画像の読み込み
img_cv = cv2.imread('apple-1834639_640.jpg')
img_cv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB)
mask_cv = cv2.imread('apple-1834639_640_mask.jpg', 0)

# 水増し処理
augmentation = albu_aug()
imgs = augmentation(image = img_cv, mask = mask_cv)

# 画像表示
imgShow([imgs['image'], imgs['mask']])

albumentationsは、引数に学習データと教師データを入力できるため、同じ処理を実行することができます。そのため、学習データだけ反転するということは起きません。

albumentationsの注意点として、読み込ませる画像はarray形式にする必要があります。

ほかにも様々な水増し方法があるので、ぜひ試してみてください。（参考）