[Pytorch] Segmentation task에서 Albumentation 이용하기

인공지능 학습할 때 overfitting을 막고 일반화 성능을 올리기 위한 다양한 방법들이 존재합니다.
그 중 쉽게 적용할 수 있는 방법 중 하나가 데이터 증강(Augmentation) 입니다. 
오늘은 데이터 증강 라이브러리인 Albumentation를 이용하여 Segmentation task에 적용해본 방법을 정리하려고 합니다.

Pytorch의 경우 TorchVision 라이브러리에 있는 torchvision.transforms을 이용할 수 있지만, Albumentation 만의 직관적인 사용법과 빠른 처리속도에 사용하지 않을 이유가 없습니다. 추천추천

1. Albumentation 설치하기

albumentation은 python 3.6 버전 이상을 지원합니다.
pip를 이용하여 아주 쉽게 설치할 수 있습니다.
 pip install albumentations 

2. Albumentation 사용하기

대부분의 내용은 Albumentation documentation을 참고하였습니다.
👉 Albumentation documentation 사이트

Albumentation을 이용하면 Segmentation 학습을 위한 augmentation은 아주 간단해집니다. 
augmentation 적용 시 interpolation 함수는 cv2 라이브러리에 있는 것을 이용합니다.

Compose를 이용하여 image와 mask 각각에 augmentation을 적용하기 위한 transform 객체를 생성합니다.
원하는 augmentation 메소드들의 list를 Compose의 입력으로 넣어주면 됩니다.

import albumentations as A
import cv2

transform = A.Compose([
	A.RandomCrop(width=256, height=256),
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
    A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=10,
                        border_mode=cv2.INTER_CUBIC, p=0.5)
])

Albumentation에서 제공하는 augmentation 기법들은 contrast, resize, zoom, flip, cutout 등등.. 아주아주 많습니다.
자세한 내용은 아래 링크를 참고하시면 좋습니다.
👉 Albumentation에서 제공하는 augmentation 기법 종류

생성한 transform 객체를 이용하여 원하는 image와 mask에 augmentation을 수행합니다.
전달받은 객체에 키값으로 'image'와 'mask'를 넣어주어 augmentation이 적용된 image와 mask를 얻을 수 있습니다.

transformed = transform(image=target_image, mask=target_mask)
transformed_image = transformed['image']
transformed_mask = transformed['mask']

multi-mask인 경우, transform argument인 masks에 원하는 mask list를 넣어주면 됩니다.

masks = [mask_1, mask_2, mask_3]

transformed = transform(image=image, masks=masks)
transformed_masks = transformed['masks']

mask1 = transformed_masks[0]