Computer vision, Machine Learning, Python

Увеличиваем датасет с помощью Albumentations для CV

Время прочтения: 4 мин.

Думаю, каждый кто работал с computer vision сталкивался с тем, что его датасет очень мал и было бы неплохо расширить его с помощью модификации уже готовых данных. Например, добавить картинки, которые зеркально отражены или повернуты на некоторый градус. Это один из базовых подходов к пополнению датасета. Сегодня рассмотрим библиотеку для аугментаций. Albumentations — это Python-библиотека, с помощью которой можно расширять датасет для компьютерного зрения. Библиотека появилась благодаря решениям с Kaggle (ее разработчиками стали Kaggle Grandmaster и три Kaggle Master) и набрала 8.6k звезд на гитхабе.

Импортируем нужные библиотеки:

import cv2 
import albumentations as A
import random
from google.colab.patches import cv2_imshow

Модуль random используем для воспроизводимости полученных результатов. Стоит отметить, что библиотека предоставляет удобный интерфейс для работы с разными типами входных данных.

Для начала инициализируем пайплайн:

transform = A.Compose([
    A.RandomCrop(width=512, height=512),
    A.HorizontalFlip(p=0.5)
])

С помощью него мы можем вырезать из изображения часть 512*512. Это важно, допустим, для задачи нахождения лиц на изображениях, так как они могут быть обрезаны, поэтому данная аугментация будет довольно полезной.

image = cv2.imread('nta.jpg')

random.seed(17)
augmented = transform(image=image)
augmented_image = augmented['image']

cv2_imshow(augmented_image);

Посмотрим на результат:

Также можем вырезать часть 300*300 из исходного изображения с разными значениями так называемого “зерна”:

transform = A.Compose([
    A.RandomCrop(width=300, height=300),
    A.HorizontalFlip(p=0.5)
])

for i in [1, 3, 6]:
  image = cv2.imread('nta.jpg')
  
  random.seed(i)
  augmented = transform(image=image)
  augmented_image = augmented['image']
  
  cv2_imshow(augmented_image);

Аналогичным образом можем повернуть изображение по горизонтали:

transform = A.HorizontalFlip(p=1)

image = cv2.imread('nta.jpg')
random.seed(17)
 
augmented = transform(image = image)
augmented_image = augmented['image']
 
cv2_imshow(augmented_image);

Теперь немного усложним задачу и получим несколько повернутых на некоторое значение градусов:

transform = A.Compose([
    A.CLAHE(),
    A.RandomRotate90(),
    A.Transpose(),
    A.Cutout(num_holes=1, max_h_size=16,max_w_size = 16,p=1),
    A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.50, rotate_limit=45, p=.75),
    A.OpticalDistortion(),
])

image = cv2.imread('nta.jpg')
random.seed(17)

augmented = transform(image=image)
augmented_image = augmented['image']

cv2_imshow(augmented_image);

Данная библиотека позволяет делать не только “геометрические” аугментации (повороты, сдвиги, отражения и т.д.), но и цветовые. Попробуем изменить насыщенность изображения:

transform = A.Compose([
 A.HueSaturationValue(p=5),
])

image = cv2.imread('nta.jpg')
random.seed(17)

augmented = transform(image=image)
augmented_image = augmented['image']

cv2_imshow(augmented_image);

Или его яркость:

transform = A.Compose([
    A.RandomBrightnessContrast(p=1),    
    A.CLAHE(p=1),    
], p=1)

image = cv2.imread('nta.jpg')
random.seed(17)
 
augmented = transform(image = image)
augmented_image = augmented['image']
 
cv2_imshow(augmented_image);

В CV задачах, связанных с изображениями городской среды, можно использовать аугментации погоды, будь то дождь или солнечные блики:

transform_rain = A.Compose(
    [A.RandomRain(blur_value=1, p=1)],
)

transform_sun = A.Compose(
    [A.RandomSunFlare(p=1)],
)

random.seed(17)

augmented_sun = transform_sun(image = image)
augmented_image_sun = augmented_sun['image']
 
augmented_rain = transform_rain(image = image)
augmented_image_rain = augmented_rain['image']

cv2_imshow(augmented_image_rain);
cv2_imshow(augmented_image_sun);

Стоит отметить, что аугментация является важной частью в обучении моделей, ведь собирать данные довольно дорогая и тяжелая задача, новые данные будут требовать дополнительной разметки, а аугментации с помощью данной библиотеки расширяют датасет и добавляют туда уже размеченные данные.

Советуем почитать