Albumentations 라이브러리를 이용하여 데이터 증강하기

이미지 증강이 이미지 분류 모델에서 눈에 띄는 성능 향상을 보여주기에 언제나 사용해왔지만, 기존에 사용하던 방식인 Keras의 ImageDataGenator는 이미지의 크기가 너무 큰 경우 모든 파일을 전처리하며 가져오기에 시간이 너무 오래 걸렸다.

따라서 이번에는 내가 아는한 가장 빠른 이미지 증강 라이브러리인 Albumentation을 이용하여 이미지를 증강시키는 예제를 포스팅하겠다.

 

Albumentations를 이용하여 이미지 증강하기

 

먼저 증강시킬 이미지와 필요한 라이브러리를 가져옵니다.

 

Albumentations는 아래와 같이 Compose안에 원하는 증강법을 넣어주어 이미지 변환을 수행합니다.

 

 

증강에 확률을 설정하여 이미지를 반복해서 넣어주어 하나의 이미지와 변환법으로 모두 결과가 다른 여러개의 데이터를 만들 수 있습니다.

또한 이미지 반환 시 dictionary형으로 반환하므로 반환값의 ["image"]에서 이미지 픽셀 값만을 가져와야 합니다.

 

 

아래와 같은 코드를 이용하여 이미지를 반복하여 다른 이미지로 만들고 그에 맞게 라벨값을 추가하여 X_train과 y_train을 넣을 시 10배 크기의 X_train과 y_train을 넘파이 어레이로 반환할 수 있습니다.

 

전체 코드

 

In [2]:

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tensorflow.keras as keras
import joblib
from sklearn.preprocessing import OneHotEncoder
from sklearn.model_selection import train_test_split
import albumentations as A
from sklearn.model_selection import StratifiedKFold

In [2]:

X = joblib.load("converted_img.pkl")
y = joblib.load("label.pkl")

 

To pass an image to the augmentation pipeline you need to call the transform function created by a call to A.Compose at Step 2. In the image argument to that function, you need to pass an image that you want to augment.

In [5]:

encoder = OneHotEncoder(sparse=False)
encoder.fit(y.reshape(-1, 1))

Out[5]:

OneHotEncoder(sparse=False)

In [6]:

X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y, shuffle=True)
y_train = encoder.transform(y_train.reshape(-1, 1))
y_test = encoder.transform(y_test.reshape(-1, 1))

test_trainsform = A.Compose([
A.Resize(32, 32)
])

img_list=[]
for i in range(len(X_test)):
    img_list.append(test_trainsform(image=X_test[i])["image"])
img_list = np.float32(img_list)
X_test = img_list

In [7]:

transform_crop = A.Compose([
    A.Resize(40, 40),
    A.RandomCrop(32, 32),
    A.OneOf([
    A.HorizontalFlip(p=1), 
    A.VerticalFlip(p=1)],p=1)
])

In [8]:

transform_flip_bright = A.Compose([
    A.Resize(32, 32),
    A.HorizontalFlip(p=0.5), 
    A.VerticalFlip(p=0.5),
    A.RandomBrightness(p=0.5, limit=(-0.25, 0.25))
])

 

C:\Users\hjhhi\anaconda3\envs\python3.9.0\lib\site-packages\albumentations\augmentations\transforms.py:1800: FutureWarning: This class has been deprecated. Please use RandomBrightnessContrast
  warnings.warn(

In [9]:

transform_fiip = A.Compose([
    A.Resize(32, 32),
    A.HorizontalFlip(p=0.5), 
    A.VerticalFlip(p=0.5),
])

In [10]:

temp=[]
for _ in range(8):
    temp.append(transform_crop(image=X_train[0])["image"])

In [11]:

fig = plt.figure(figsize=(14, 10))
for i in range(1, 9):
    ax = plt.subplot(2, 4, i)
    ax.axes.xaxis.set_visible(False)
    ax.axes.yaxis.set_visible(False)
    plt.imshow(temp[i-1])

 

In [12]:

temp=[]
for _ in range(8):
    temp.append(transform_fiip(image=X_train[0])["image"])

In [13]:

fig = plt.figure(figsize=(14, 10))
for i in range(1, 9):
    ax = plt.subplot(2, 4, i)
    ax.axes.xaxis.set_visible(False)
    ax.axes.yaxis.set_visible(False)
    plt.imshow(temp[i-1])

 

In [14]:

temp=[]
for _ in range(8):
    temp.append(transform_flip_bright(image=X_train[0])["image"])

In [15]:

fig = plt.figure(figsize=(14, 10))
for i in range(1, 9):
    ax = plt.subplot(2, 4, i)
    ax.axes.xaxis.set_visible(False)
    ax.axes.yaxis.set_visible(False)
    plt.imshow(temp[i-1])

 

In [18]:

def make_train_data(augment):
    img_list = []
    label_list = []
    for _ in range(10):
        for i in range(len(X_train)):
            img_list.append(augment(image=X_train[i])["image"])
            label_list.append(y_train[i])
    img_list = np.float32(img_list)
    img_list = img_list.reshape(-1, 32, 32, 3)
    label_list = np.float32(label_list)
    return [img_list, label_list]

 

---

Albumentations 페이지: https://albumentations.ai/

Albumentations

코드 링크: https://github.com/hykhhijk/Rock-classifier/blob/master/Augmentation.ipynb

GitHub - hykhhijk/Rock-classifier: Rock classifier using CNN