セマンティックセグメンテーションを試してみた
勉強のためセマンティックセグメンテーションをやってみましたのでまとめます。
セマンティック セグメンテーション (Semantic Segmentation) は、画像内の全画素にラベルやカテゴリを関連付けるディープラーニング (Deep Learning) のアルゴリズムです。
Pytorchベースのセマンティックセグメンテーション用の便利なライブラリ、segmentation_models_pytorchを使っています。
データセットはMIT Scene Parsing Benchmarkを使っています。
コードはおおむねsegmentation_models_pytorchのサンプルコードを使用しています。
Google Colab で実行できるようにしていますので、こちらをご参照ください。
なお、データセットはマウントしたGoogle Drive上に保存して使っています。
コード中の xxx の部分をご自身の環境に合わせて書き換えて使用ください。
import os
os.chdir('/content/drive/MyDrive/xxxxxx')
慣れないため不備もあろうかと思いますが、温かい目で指摘のほどお願いいたします。
主な変更点のみ記載し、詳細は割愛いたします。
訓練用データセットが約2万と大きいため、64x64にリサイズにするよう設定しています。
numpyのarrayデータのままだとリサイズできないので、下記の回答に習って画像に変換してからリサイズしています。
なぜか 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回では微妙な感じでしょうか・・・。
以上になります、最後までお読みいただきありがとうございました。
Discussion