[Python] albumentations으로 라이브러리로 augmentation 적용하기 (feat. OneOf 활용, 여러 이미지 적용)

일반적으로 딥러닝 모델은 많은 이미지를 학습할수록 더 좋은 성능을 내지만 실질적으로 이미지의 수가 부족한 경우가 매우 많다. 그럴 때 사용하는 것이 image augmentation으로 가지고 있는 image를 변형하여 모델이 좀 더 다양한 데이터를 학습할 수 있도록 한다. 이런 image augmentation의 적용을 쉽게하고자 albumentations라는 라이브러리가 존재한다. 이 라이브러리를 더욱 확장성 있게 활용하는 방법(여러 이미지 적용, OneOf 활용 등)들을 알아보자.

 

출처 : (https://albumentations.ai/)

 

기본 

가장 간단한 활용할 수 있는 코드를 살펴보자. 먼저 augmentation 정의하는 코드는 다음과 같다.

from albumentations import *

def test_aug():
    ret = Compose([
        ElasticTransform(),
        GridDistortion(),
        HorizontalFlip(),
        OpticalDistortion(),
        ShiftScaleRotate(),
        VerticalFlip()
    ])
    return ret

이제 이미지에 대해서 적용하려면

import cv2

image1 = cv2.cvtColor(cv2.imread('../sample.jpg'), cv2.COLOR_BGR2RGB)[:,100:500]
mask = np.zeros((512, 512))
mask = cv2.fillPoly(mask, pts1, 1)
mask = cv2.fillPoly(mask, pts2, 2)
mask = cv2.fillPoly(mask, pts3, 3)

plt.figure(figsize=(10,10))
plt.subplot(2,2,1)
plt.imshow(image1)
plt.title('origin_1')
plt.subplot(2,2,2)
plt.imshow(mask, vmin=0, vmax=3, cmap='Oranges')
plt.title('origin_mask')

# 적용하는 부분
test_a = test_aug()
test = test_a(image=image1, mask=mask)

plt.subplot(2,2,1)
plt.imshow(test['image'])
plt.title('test_1')
plt.subplot(2,2,2)
plt.imshow(test['mask'], vmin=0, vmax=3, cmap='Oranges')
plt.title('test1_mask')

 

다음과 같다.

 

여기서 사용한 이미지 처리는 정확한 의미는 아니지만 이미지를 회전하는 정도의 처리 3개(HorizontalFlip(), VerticalFlip(), ShiftScaleRotate())와 이미지를 변형하는 정도의 처리 3개(ElasticTransform(), GridDistortion(), OpticalDistortion())라고 보자. 각각의 처리에는 p(확률값)을 주어서 처리할 수 있는데 디폴트 값은 0.5이다. 따라서 위의 방식대로 하면 6개의 처리가 각각 0.5의 확률로 처리가 된다. 확률적으로 정확하게 계산해봐야겠지만 64(= $2^6$)가지의 경우의 수가 나오나 직관적으로 기본이미지에 대해 확률처리가 되므로 변형된 것이 적게 나오는 아쉬움이 있을 것 같다. 또한, 너무 다양한 처리를 하여 오히려 모델에 안 좋은 효과를 줄 수 있을 것이다.

 

 

OneOf 활용

이를 고려하면 (이미지를 회전하는 정도의 처리 3개 + 기본이미지)와 (이미지를 변형하는 정도의 처리 3개 + 기본이미지)에서 각각 한 개의 처리만 동일한 확률로 나오게 하여 경우의 수를 고려해주는 방법으로 처리해보자. 이를 위해서는 OneOf를 활용하면 된다.

 

def test1_aug():
    ret = Compose([
        OneOf([
        HorizontalFlip(p=1),
        VerticalFlip(p=1),
        ShiftScaleRotate(p=1)
        ],p=0.75),
        OneOf([
        ElasticTransform(p=1, alpha=1, sigma=20, alpha_affine=15, interpolation=1, border_mode=1),
        GridDistortion(p=1),
        OpticalDistortion(p=1)
        ],p=0.75)
    ])
    return ret

 

OneOf 함수에 3개씩 담아준 후 안에 처리는 확률을 1로 OneOf의 확률을 0.75로 하면 될 것이다. (저의 단순한 계산이므로 잘못될 시 지적부탁드립니다.) 

 

이제 이미지를 2개라고 해보자. 마스크를 포함하여 총 3개에 대해 augmentation을 하려면 어떻게 해야할까?

image2 = cv2.resize(image1[100:300,100:300], (400,400), interpolation=cv2.INTER_CUBIC)

plt.figure(figsize=(15,5))
plt.subplot(1,3,1)
plt.imshow(image1)
plt.title('origin_1')
plt.subplot(1,3,2)
plt.imshow(image2)
plt.title('origin_2')
plt.subplot(1,3,3)
plt.imshow(mask, vmin=0, vmax=3, cmap='Oranges')
plt.title('origin_mask')

 

additional_tagets={}의 활용

이를 위해 additional_targets={'image1':'image'}를 추가해주자.

def test2_aug():
    ret = Compose(
        [
            OneOf([
                HorizontalFlip(p=1),
                VerticalFlip(p=1),
                ShiftScaleRotate(p=1)
            ],p=0.75),
            OneOf([
                ElasticTransform(p=1, alpha=1, sigma=20, alpha_affine=15, interpolation=1, border_mode=1),
                GridDistortion(p=1),
                OpticalDistortion(p=1)
            ],p=0.75)
        ],
        additional_targets={'image1':'image'}
    )
    return ret

이를 총 3개의 이미지에 대해 적용하면

test2_a = test2_aug()
test2 = test2_a(image=image1, image1=image2, mask=mask)

여러 이미지에 대해 적용할 수 있다. 마찬가지로 mask를 추가할 수도 있고 bounding bax나 keypoint에 대해서도 적용할 수 있을 것이다.

 

 

-히비스서커스-

 

참고자료

https://hoya012.github.io/blog/albumentation_tutorial/

albumentations - fast image augmentation library 소개 및 사용법 Tutorial

https://albumentations.ai/docs/examples/example_multi_target/

Albumentations Documentation - Applying the same augmentation with the same parameters to multiple images, masks, bounding boxes