MMsegmentation Augmentation 사용법

[kimkihoon0515]

다들 MMObjectDetection때 해보셔서 아시리라고 생각되지만 혹시나 그래도 필요하신분을 위해 작성했습니다.
Config에서 Dataset을 적어놓은 파이썬 파일을 보시면 저런 pipeline들이 있습니다.

train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations'),
    dict(type='Resize', img_scale=(512, 2048), ratio_range=(0.5, 2.0)),
    dict(type='CLAHE'),
    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
    dict(type='RandomFlip', prob=0.5),
    dict(type='RandomRotate',prob=0.4,degree=30),
    dict(type='PhotoMetricDistortion'), # random brightness, random contrast, convert color from BGR to HSV, random saturation, random hue, convert color from HSV to BGR
    dict(type='Normalize', **img_norm_cfg),
    dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_semantic_seg']),
]
valid_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=[(512, 512), (768,768), (1024, 1024), (1280, 1280), (1536, 1536)],
        flip=True,
        flip_direction=['horizontal', 'vertical'],
        transforms=[
            dict(type='Resize', keep_ratio=True),
            dict(type='RandomFlip'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img']),
        ])
]

이것들은 통해 augmentation을 dict()형태로 적용이 가능한데 albumentation을 따로 할 필요없이 거의 모든 augmentation들이 이미 기능이 구현이 되어있습니다.
/opt/ml/mmsegmentation/mmseg/datasets/pipelines/transforms.py
위 경로로 가시면 class별로 Augmentation들이 나눠져있습니다. 예를 들어 RandomFlip의 경우에는 다음과 같이 선언되어 있습니다.

class RandomFlip(object): # dict에 type=에 해당함 dict(type='RandomFlip')
    """Flip the image & seg.

    If the input dict contains the key "flip", then the flag will be used,
    otherwise it will be randomly decided by a ratio specified in the init
    method.

    Args:
        prob (float, optional): The flipping probability. Default: None.
        direction(str, optional): The flipping direction. Options are
            'horizontal' and 'vertical'. Default: 'horizontal'.
    """

    @deprecated_api_warning({'flip_ratio': 'prob'}, cls_name='RandomFlip')
    def __init__(self, prob=None, direction='horizontal'):
        self.prob = prob
        self.direction = direction
        if prob is not None:
            assert prob >= 0 and prob <= 1
        assert direction in ['horizontal', 'vertical']

    def __call__(self, results):
        """Call function to flip bounding boxes, masks, semantic segmentation
        maps.

        Args:
            results (dict): Result dict from loading pipeline.

        Returns:
            dict: Flipped results, 'flip', 'flip_direction' keys are added into
                result dict.
        """

        if 'flip' not in results:
            flip = True if np.random.rand() < self.prob else False
            results['flip'] = flip
        if 'flip_direction' not in results:
            results['flip_direction'] = self.direction
        if results['flip']:
            # flip image
            results['img'] = mmcv.imflip(
                results['img'], direction=results['flip_direction'])

            # flip segs
            for key in results.get('seg_fields', []):
                # use copy() to make numpy stride positive
                results[key] = mmcv.imflip(
                    results[key], direction=results['flip_direction']).copy()
        return results

    def __repr__(self):
        return self.__class__.__name__ + f'(prob={self.prob})' # dict에서 prob 등 꼭 필요한 인자들로 구성되어있습니다.

그리고 잘 보시면 docstring에 어떤 인자가 어떤 역할을 하는지 어떤 식으로 augmentation이 적용되는지 설명이 잘 되어있습니다.
잘 참고하셔서 적용하시면 될 것 같습니다. 다만 RandomMosaic는 적용이 안되는데 Dataset을 조금 수정을 해주어야 적용이 가능해보입니다.