mmsegmentationでカスタムデータセットを使用して学習をする方法

概要

mmsegmentationとは、OpenMMLabプロジェクトの一つであるPytorchをベースとしたオープンソースセグメンテーションツールボックスです。この記事ではmmsegmentationでカスタムデータセットを使用して学習と評価をする方法を説明します。

新しいカスタムデータセットを登録する

まず新しいカスタムデータセットをmmsegmentationに登録します。

mmseg/datasets/example.py に新しいファイルを作成します。exampleは適宜変更してください。

from mmseg.registry import DATASETS
from .basesegdataset import BaseSegDataset


@DATASETS.register_module()
class ExampleDataset(BaseSegDataset):

    METAINFO = dict(
        classes=('xxx', 'xxx', ...), #クラス名の一覧を入力してください
        palette=[[x, x, x], [x, x, x], ...]) #クラスに対応する色情報（RGB）を入力してください

    def __init__(self, arg1, arg2):
        pass

mmseg/datasets/__init__.py にモジュールをインポートします。

from .example import ExampleDataset

configs/_base_/datasets/example_dataset.py に新しいデータセット設定ファイルを作成して使用します。このとき、img_pathなどのpathも合わせることを忘れないでください。

dataset_type = 'ExampleDataset'
data_root = 'data/example/'
...

mmseg/utils/class_names.py にデータセットのメタ情報を追加します。

def example_classes():
    return [
        'xxx', 'xxx',
        ...
    ]

def example_palette():
    return [
        [x, x, x], [x, x, x],
        ...
    ]
dataset_aliases ={
    'example': ['example', ...],
    ...
}

configファイルを作成する

次にconfigファイルを作成します。

例として、deeplabv3+を使用した場合、configs/example/deeplabv3plus_r50-d8_4xb4-40k_example-512x512.py というファイルを作成します。

_base_ = [
    '../_base_/models/deeplabv3plus_r50-d8.py',
    '../_base_/datasets/example_dataset.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_20k.py'
]
crop_size = (512, 512)
data_preprocessor = dict(size=crop_size)
model = dict(
    data_preprocessor=data_preprocessor,
    decode_head=dict(num_classes=2),
    auxiliary_head=dict(num_classes=2))

※configファイルの命名規則については、mmsegmentationのチュートリアルのConfig Name Styleを参照してください。

学習を開始する

最後にカスタムデータセットでの学習を開始します。

mmsegmentation下で次のコマンドを打ちます。

$ python tools/train.py configs/original/deeplabv3plus_r50-d8_4xb4-40k_example-512x512.py --work-dir ${WORK_DIR}

学習が開始されます。

まとめ

今回は、mmsegmentationでカスタムデータセットを使用した学習までを実際に行ってみました。また、configファイルを編集することで希望のモデルやパラメータで学習を行うことができます。

参考

mmsegmentation

備考

画像解析を用いたプロダクト開発をご検討中ですか？作業方法や具体作業のご発注等について、より詳細なご説明やご相談をご希望の場合は、以下フォームよりお気軽にお問い合わせくださいませ。