Время прочтения: 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);
Стоит отметить, что аугментация является важной частью в обучении моделей, ведь собирать данные довольно дорогая и тяжелая задача, новые данные будут требовать дополнительной разметки, а аугментации с помощью данной библиотеки расширяют датасет и добавляют туда уже размеченные данные.