セマンティックセグメンテーションを試してみた

3 min read読了の目安(約2900字

勉強のためセマンティックセグメンテーションをやってみましたのでまとめます。

セマンティック セグメンテーション (Semantic Segmentation) は、画像内の全画素にラベルやカテゴリを関連付けるディープラーニング (Deep Learning) のアルゴリズムです。

https://jp.mathworks.com/solutions/image-video-processing/semantic-segmentation.html

Pytorchベースのセマンティックセグメンテーション用の便利なライブラリ、segmentation_models_pytorchを使っています。

データセットはMIT Scene Parsing Benchmarkを使っています。

コードはおおむねsegmentation_models_pytorchのサンプルコードを使用しています。

Google Colab で実行できるようにしていますので、こちらをご参照ください。

Open In Colab

なお、データセットはマウントしたGoogle Drive上に保存して使っています。

コード中の xxx の部分をご自身の環境に合わせて書き換えて使用ください。

import os
os.chdir('/content/drive/MyDrive/xxxxxx')

慣れないため不備もあろうかと思いますが、温かい目で指摘のほどお願いいたします。

主な変更点のみ記載し、詳細は割愛いたします。

訓練用データセットが約2万と大きいため、64x64にリサイズにするよう設定しています。

numpyのarrayデータのままだとリサイズできないので、下記の回答に習って画像に変換してからリサイズしています。

https://teratail.com/questions/152443

なぜか Pytorch や segmentation_models_pytorch の Resize はうまくいきませんでした...。

リサイズする画像の大きさ

IMAGE_SIZE = 64

データローダー

class Dataset(BaseDataset):
~省略~
    def __getitem__(self, i):
~省略~
        image = Image.fromarray(np.uint8(image))
        image = np.asarray(image.resize((IMAGE_SIZE, IMAGE_SIZE)))
~省略~
        mask = Image.fromarray(np.uint8(mask))
        mask = np.asarray(mask.resize((IMAGE_SIZE, IMAGE_SIZE))) 
~省略~

リサイズした画像を出力すると以下のようになります。

さすがにちょっと荒いですね...。

データ拡張するとこんな感じになりました。

元々のサンプルコードのモデルは ResNext50 でしたが、少しでも軽くしたくて VGG16 にしています。

バッチサイズは学習用の画像が20,210枚なので、20,210の約数である47にしています。

なお、バリデーションデータは2,000枚でした。

※ 割り切れないバッチサイズを設定するとエラーになりました。

モデル作成と学習

パラメーター

ENCODER = 'vgg16'

BATCH_TRAIN = 47
BATCH_VAL = 50

元のサンプルコードでは推論時にテスト用の画像とマスク画像を使っていますが、自分が使用したデータセットにはテスト用のマスク画像がなかったので、バリデーションデータの画像とマスク画像を使いました。

推論

テスト用のデータセットの作成

test_dataset = Dataset(
    x_valid_dir, 
    y_valid_dir,
    augmentation=get_validation_augmentation(), 
    preprocessing=get_preprocessing(preprocessing_fn),
)

test_dataloader = DataLoader(test_dataset)

推論結果は以下のとおりです。

一番右が推論結果になります。

エポック数10回では微妙な感じでしょうか・・・。

以上になります、最後までお読みいただきありがとうございました。