MMSeg에서 새로운 증강 기법 만들기 | Custom Transforms

MMSegmentation에 추가하기

1. 클래스 생성 

MMCV에서 정의된 여러 변형 기법들이 있지만[1], 사용하기에 적합하지 않거나 새로운 기법을 적용하고 싶을 수 있습니다. 이럴 경우 새로운 사용자 정의 변형(transformation) 기법을 만들어 등록하는 과정을 거칩니다. 아래 예시는 프로젝트에서 utils라는 하위 디렉토리를 만들고 그 안에 function.py를 만든 후 RandomSharpen이라는 transformation 클래스를 만들었습니다. 

 

from mmseg.registry import DATASETS
from .basesegdataset import BaseSegDataset
from PIL import ImageEnhance, Image


@TRANSFORMS.register_module()
class RandomSharpen(BaseTransform):
    def __init__(self, prob=0.5, sharpness_factor=(0.8, 1.2)):
        self.prob = prob
        self.sharpness_factor = sharpness_factor

    def _apply_sharpen(self, img):
        enhancer = ImageEnhance.Sharpness(img)
        factor = np.random.uniform(self.sharpness_factor[0], self.sharpness_factor[1])
        return enhancer.enhance(factor)

    def transform(self, results):
        if np.random.rand() < self.prob:
            img = results['img']
            img_pil = Image.fromarray(img)  # numpy -> PIL
            img_sharpened = self._apply_sharpen(img_pil)
            results['img'] = np.array(img_sharpened)  # PIL -> numpy
        return results

 

새로운 transformation 클래스를 만들 때는 BaseTransform 을 상속받아야만 합니다. 그리고 transform 메서드를 만들어야 하는데, 자세한 내용은 참고 문헌을 참고해주시기 바랍니다. [2]

 

2. 불러오기

새롭게 정의된 클래스를 아래와 같이 가져올 수 있습니다. 

 

from .utils.function import RandomSharpen

 

또는 custom_imports를 통해 불러올 수 있습니다.

 

custom_imports = dict(
    imports=[
        'utils.function'
    ],  
    allow_failed_imports=False 
)

3. 적용하기

import 로 가져온 클래스는 아래 처럼 사용할 수 있습니다.

 

transform = RandomSharpen()
data_dict = {'img': np.random.rand(224, 224, 3)}
data_dict = transform(data_dict)
processed_img = data_dict['img']

 

또는 config를 통해서도 아래와 같이 사용할 수 있습니다.

 

pipeline = [
    ...
    dict(type='RandomSharpen'),
    ...
]

참고자료

[1] https://mmcv.readthedocs.io/en/latest/understand_mmcv/data_transform.html

[2] https://mmcv.readthedocs.io/en/latest/understand_mmcv/data_transform.html#design-of-data-transformation

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