[Augmentation] mmseg Albumentation 추가

[baekkr95]

Background

• mmseg에 Albumentation를 추가하는 코드
• mmsegmenation 파일은 gitignore되어 있기 때문에 issue로 코드를 올렸습니다. 꼭 확인 부탁드려요!

Content

• 레퍼런스 : https://github.com/open-mmlab/mmsegmentation/pull/392/files#diff-1be06a265a3c5d0a62c8b24b9c75cfcb3f386ec0087ea0f4a2ff10ff0ab7d3ba
• 파일경로 : mmsegmentation > mmseg > datasets > pipelines > transforms.py 파일 안에 Albumentation 코드를 추가
• pipelines > __init__.py에 Albu 클래스를 import를 하면 됨.

Details

아래 코드를 transforms.py 에 추가하면 됩니다. 이후 config datasets에서 mmdetection에서 albumentation 사용했던 것처럼 사용하면 됩니다.

import copy
import mmcv
import numpy as np
from mmcv.utils import deprecated_api_warning, is_tuple_of
from numpy import random

from ..builder import PIPELINES

try:
    import albumentations
    from albumentations import Compose
except ImportError:
    albumentations = None
    Compose = None

@PIPELINES.register_module()
class Albu(object):
    """Albumentation augmentation. Adds custom transformations from
    Albumentations library. Please, visit
    `https://albumentations.readthedocs.io` to get more information. An example
    of ``transforms`` is as followed:
    .. code-block::
        [
            dict(
                type='ShiftScaleRotate',
                shift_limit=0.0625,
                scale_limit=0.0,
                rotate_limit=0,
                interpolation=1,
                p=0.5),
            dict(
                type='RandomBrightnessContrast',
                brightness_limit=[0.1, 0.3],
                contrast_limit=[0.1, 0.3],
                p=0.2),
            dict(type='ChannelShuffle', p=0.1),
            dict(
                type='OneOf',
                transforms=[
                    dict(type='Blur', blur_limit=3, p=1.0),
                    dict(type='MedianBlur', blur_limit=3, p=1.0)
                ],
                p=0.1),
        ]
    Args:
        transforms (list[dict]): A list of albu transformations
        keymap (dict): Contains {'input key':'albumentation-style key'}
    """
    def __init__(self, transforms, keymap=None, update_pad_shape=False):
        # Args will be modified later, copying it will be safer
        transforms = copy.deepcopy(transforms)
        if keymap is not None:
            keymap = copy.deepcopy(keymap)
        self.transforms = transforms
        self.filter_lost_elements = False
        self.update_pad_shape = update_pad_shape
        self.aug = Compose([self.albu_builder(t) for t in self.transforms])
        if not keymap:
            self.keymap_to_albu = {'img': 'image', 'gt_semantic_seg': 'mask'}
        else:
            self.keymap_to_albu = keymap
        self.keymap_back = {v: k for k, v in self.keymap_to_albu.items()}
    def albu_builder(self, cfg):
        """Import a module from albumentations.
        It inherits some of :func:`build_from_cfg` logic.
        Args:
            cfg (dict): Config dict. It should at least contain the key "type".
        Returns:
            obj: The constructed object.
        """
        assert isinstance(cfg, dict) and 'type' in cfg
        args = cfg.copy()
        obj_type = args.pop('type')
        if mmcv.is_str(obj_type):
            obj_cls = getattr(albumentations, obj_type)
        else:
            raise TypeError(f'type must be str, but got {type(obj_type)}')
        if 'transforms' in args:
            args['transforms'] = [
                self.albu_builder(transform)
                for transform in args['transforms']
            ]
        return obj_cls(**args)
    @staticmethod
    def mapper(d, keymap):
        """Dictionary mapper.
        Renames keys according to keymap provided.
        Args:
            d (dict): old dict
            keymap (dict): {'old_key':'new_key'}
        Returns:
            dict: new dict.
        """
        updated_dict = {}
        for k, v in zip(d.keys(), d.values()):
            new_k = keymap.get(k, k)
            updated_dict[new_k] = d[k]
        return updated_dict
    def __call__(self, results):
        # dict to albumentations format
        results = self.mapper(results, self.keymap_to_albu)
                
        results = self.aug(**results)
            
        # back to the original format
        results = self.mapper(results, self.keymap_back)

        # update final shape
        if self.update_pad_shape:
            results['pad_shape'] = results['img'].shape

        return results

    def __repr__(self):
        repr_str = self.__class__.__name__ + f'(transforms={self.transforms})'
        return repr_str

[seonahmin]

잘 적용됐는지 train_data를 시각화하는 부분은 AI stages에서 조교님이 올려주신 방법을 사용해서 확인할 수 있을 것 같습니다.

[baekkr95]

• 선아님 말씀대로 시각화해본 결과 적용이 잘 되고 있습니다.
• 왼쪽이 오리지널, 오른쪽이 albumentation shiftscalerotate 적용한 사진입니다.