Albumentation 사용해서 Augmentation하기

예전에 Pytorch의 Image augmentation 방법을 보고 커스텀이 되게 편리하게 구성된 것 같아 TensorFlow도 이와 같은 방법으로 augmentation하도록 만들어줬으면 좋겠다는 생각을 하고 있었는데, 연동 가능한 라이브러리 albumentation이 있었습니다.
(만들어진지 좀 됬는데 늦게 알게 됨..)

transforms.Compose([
	transforms.CenterCrop(10),
	transforms.ToTensor(),
])

torchvision의 augmentation 방법

사실 아래 공식 홈페이지 튜토리얼 코드를 보면 TensorFlow도 이와 같은 방법으로 Keras Layer를 활용한 augmentation이 가능토록 제공하고 있고, 앞으로도 이러한 방법으로 제공할 예정인가 싶습니다.

data_augmentation = tf.keras.Sequential([
  layers.experimental.preprocessing.RandomFlip("horizontal_and_vertical"),
  layers.experimental.preprocessing.RandomRotation(0.2),
])

출처: www.tensorflow.org/tutorials/images/data_augmentation?hl=ko

---

tf.image를 사용하는 방법

albumentation 라이브러리를 사용해보기 전에, tf.image를 사용해서 어떻게 augmentation 할 수 있는지 코드를 첨부합니다.

test = tf.data.Dataset.from_tensor_slices(dict(df))

# 이미지와 레이블을 얻습니다.
def get_image_label(dt):
    img_path = dt['image']
    
    image = tf.io.read_file(img_path)
    image = tf.image.decode_jpeg(image)
    image = tf.image.convert_image_dtype(image, tf.float32)
    image = tf.image.resize(image, [IMG_SIZE, IMG_SIZE])
    image = (image / 255.0)

    label = []
    
    for key in class_col:
        label.append(dt[key])
    
    return image, label

# data augment
# refer: https://www.tensorflow.org/tutorials/images/data_augmentation
def augment(image,label):
    # Add 6 pixels of padding
    image = tf.image.resize_with_crop_or_pad(image, IMG_SIZE + 6, IMG_SIZE + 6) 
    # Random crop back to the original size
    image = tf.image.random_crop(image, size=[IMG_SIZE, IMG_SIZE, 3])
    image = tf.image.random_brightness(image, max_delta=0.5) # Random brightness
    image = tf.clip_by_value(image, 0, 1)

    return image, label

    
dataset = test.map(get_image_label)
dataset = dataset.shuffle(50).map(augment).batch(4)

이 방법도 괜찮긴하지만 문제는 tf.Dataset 작동 구조를 알아야 정확히 쓸 수 있습니다.
map 함수 작동 방식이라던가, shuffle과 batch의 위치 등등..

확실히 위 방법보다 augmentation 방법을 Layer 구조로 가져가는게 훨씬 가독성이 좋아보이기도 합니다. 아래처럼요.

model = tf.keras.Sequential([
  resize_and_rescale,
  data_augmentation,
  layers.Conv2D(16, 3, padding='same', activation='relu'),
  layers.MaxPooling2D(),
  # Rest of your model
])

resize_and_rescale과 data_augmentation을 Sequential 안에 Layer 형태로 사용하고 있습니다.

하지만 아직 많이 쓰이고 있는 방법은 아닌 것 같기에 확실하게 자리잡기 전(?)까지 다른 방법을 사용해도 꽤나 무방합니다.

 

Albumentation 라이브러리를 사용하는 방법

사용하는 방법은 꽤나 단순합니다. 위에서 예제로 보여드렸던 코드와 형식이 동일하니까요.

먼저, 필요 라이브러리를 불러옵니다.

from tensorflow.keras.datasets import cifar10

import albumentations
import cv2
import numpy as np
import matplotlib.pyplot as plt
import cv2
# augmentation method를 import합니다.
from albumentations import (
    Compose, HorizontalFlip, CLAHE, HueSaturationValue,
    RandomBrightness, RandomContrast, RandomGamma,
    ToFloat, ShiftScaleRotate
)

torchvision의 Compose, TensorFlow의 Sequential과 Albumentation의 Compose의 사용되는 장소가 같습니다. 위에서 호출한 함수 외에도 다양한 augmentation 방법을 제공합니다. 더 궁금하면 공식 문서를 참조하고, 자세히 설명되어 있습니다.

albumentations.ai/docs/getting_started/image_augmentation/

이제 다양한 함수를 Compose안에 list 형태로 제공하고, 사용할 준비를 끝마칩니다.

# 각 함수에 대한 설명은
# https://albumentations.ai/docs/
# document를 참고하세요.
Aug_train = Compose([
    HorizontalFlip(p=0.5),
    RandomContrast(limit=0.2, p=0.5),
    RandomGamma(gamma_limit=(80, 120), p=0.5),
    RandomBrightness(limit=0.2, p=0.5),
    HueSaturationValue(hue_shift_limit=5, sat_shift_limit=20,
                       val_shift_limit=10, p=.9),
    ShiftScaleRotate(
        shift_limit=0.0625, scale_limit=0.1, 
        rotate_limit=15, border_mode=cv2.BORDER_REFLECT_101, p=0.8), 
    ToFloat(max_value=255)
])

Aug_test = Compose([
    ToFloat(max_value=255)
])

 

실험하기에 가장 좋은 예제는 tensorflow.keras.datasets에서 제공하는 CIFAR-10입니다.
또, Sequence 클래스를 상속받아 제네레이터처럼 활용할 수 있도록 합니다.
아래 코드 __getitem__ 부분의 self.augment(image=x)["image"]에서 변환 작업이 수행됩니다.

from tensorflow.python.keras.utils.data_utils import Sequence

# Sequence 클래스를 상속받아 generator 형태로 사용합니다.
class CIFAR10Dataset(Sequence):
    def __init__(self, x_set, y_set, batch_size, augmentations):
        self.x, self.y = x_set, y_set
        self.batch_size = batch_size
        self.augment = augmentations

    def __len__(self):
        return int(np.ceil(len(self.x) / float(self.batch_size)))

    # 지정 배치 크기만큼 데이터를 로드합니다.
    def __getitem__(self, idx):
        batch_x = self.x[idx * self.batch_size:(idx + 1) * self.batch_size]
        batch_y = self.y[idx * self.batch_size:(idx + 1) * self.batch_size]
        
        # augmentation을 적용해서 numpy array에 stack합니다.
        return np.stack([
            self.augment(image=x)["image"] for x in batch_x
        ], axis=0), np.array(batch_y)

# CIFAR-10 Dataset을 불러옵니다.
(x_train, y_train), (x_test, y_test) = cifar10.load_data()

BATCH_SIZE = 8

# Dataset을 생성합니다.
train_gen = CIFAR10Dataset(x_train, y_train, BATCH_SIZE, Aug_train)
test_gen = CIFAR10Dataset(x_test, y_test, BATCH_SIZE, Aug_test)

train_gen을 통해 이미지를 그려보면 다음과 같이 변환이 일어난 것을 볼 수 있습니다.

# 데이터를 그려봅시다.
images, labels = next(iter(train_gen))

fig = plt.figure()

for i, (image, label) in enumerate(zip(images, labels)):
    ax = fig.add_subplot(3, 3, i + 1)
    ax.imshow(image)
    ax.set_xlabel(label)
    ax.set_xticks([]); ax.set_yticks([])

plt.tight_layout()
plt.show()

Augmentation이 수행된 이미지

 

Reference

텐서플로우 공식 홈페이지

Albumentation Document

medium.com/the-artificial-impostor/custom-image-augmentation-with-keras-70595b01aeac