MMSeg에서 새로운 데이터셋 정의하기 | Custom Dataset

MMSegmentation에 추가하기

1. 데이터셋 추가 

새로운 사용자 정의 데이터셋을 mmseg/datasets/example.py로 만듭니다. 새로운 모듈을 추가하기 위해서는 registry에 정의되어 있는 클래스의 register_module을 통해 등록합니다. 여기서는 MMSegmentation에서 제공하고 있는 BaseSegDataset 클래스를 상속받아 작성하는 것을 가정했습니다. 

 

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], ...])

    def __init__(self, arg1, arg2):
        pass

2. 불러오기

새롭게 정의된 데이터셋을 mmseg/datasets/__init__.py에 아래와 같이 추가해줍니다.

 

from .example import ExampleDataset

3. config 파일에 추가

새롭게 만들어진 데이터셋에 대한 config를 configs/_base_/datasets/example_dataset.py에 아래와 같이 작성해서 추가해줍니다. 

 

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

4.정보 추가

메타 데이터(클래스, 팔레트 등)를 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', ...],
    ...
}

참고자료

[1] https://mmsegmentation.readthedocs.io/en/latest/advanced_guides/add_datasets.html