[딥러닝] Albumentations으로 증강 쉽게하기(Augmentation-python)
Albmentations
Albumentations는 이미지를 손쉽게 augmentation 해주는 python 라이브러리이다. 코드 몇 줄로 매우 쉽게 구현할 수 있다. 또한 다양한 영상변환 알고리즘을 제공하고 있고, 그 처리속도도 매우 빨라 딥러닝 전처리 용으로 유용하게 사용할 수 있다.
1. 이미지
샘플 이미지로 Pascal VOC 중 하나의 이미지를 사용했습니다.
A.Compose() 안에 augmentation하고 싶은 함수가 들어있는 리스트를 입력하면 됩니다. 각 함수들에는 p라는 파라미터를 입력받게 되어있는데, probability의 약자로 해당 함수가 적용될 확률을 입력하면 됩니다. 1이라면 무조건 적용하는 것이고, 0.5라면 절반의 확률로 적용됩니다.
이후 transform 함수에 image파라미터에 image 변수를 넣으면 dictionary 형태로 반환되는데, image라는 key의 value값으로 결과가 반환됩니다.
(사용가능한 transform 함수목록은 공식문서에서 확인할 수 있습니다.)
import albumentations as A
import matplotlib.pyplot as plt
image = plt.imread('sample.jpg')
transform = A.Compose([
# 변환할 함수들 리스트
A.HorizontalFlip(p=1),
])
transformed = transform(image=image)
transformed_image = transformed['image']
plt.imshow(transformed_image)
plt.xticks([]); plt.yticks([])
plt.show()
2. 마스크
위와 사용방법은 같고, transform 함수의 파라미터로 mask를 입력하면 됩니다. 위와 마찬가지고 dictionary형태로 반환되는데, 각각의 key를 조회하면 결과를 반환할 수 있습니다.
import albumentations as A
import matplotlib.pyplot as plt
image = plt.imread('sample.jpg')
mask = plt.imread('mask.jpg')
transform = A.Compose([
# 변환할 함수들 리스트
A.HorizontalFlip(p=1),
])
transformed = transform(image=image, mask=mask)
transformed_image = transformed['image']
transformed_mask = transformed['mask']3. Bounding Box
Bounding Box는 위와 아주 조금 다릅니다. A.Compose로 augmentation 목록을 입력할 때, bbox_params라는 파라미터를 입력해주어야 합니다.
Bounding box의 format은 Object Detection의 잘 알려진 데이터셋인 COCO, Pascal VOC 등의 format을 따라가기 때문에 다음과 같이 총 4가지를 지원합니다. 'coco', 'pascal_voc', 'albumentations', 'yolo'
각 format별로 bounding box 좌표값을 처리하는 방식이 다르니 유의해서 사용하면 됩니다.
(A.BBoxParams의 자세한 사항은 공식문서를 참조하세요.)
이후에는 위와 마찬가지로 transform 함수에 bboxes 파라미터에 bbox값들을 입력하고, 반환된 dictionary에서 bboxes라는 key로 결과를 얻어낼 수 있습니다.
반복되는 부분이라 이해할 수 있겠지만 transform 함수에는 여러가지 파라미터를 함께 넣어도 됩니다. image, mask, bboxes, keypoints 등 augmentation을 원하는 모든 값을 입력할 수 있습니다.
아래 코드는 augmentation을 위해 반환된 bbox를 영상에 출력하는 코드입니다.
import albumentations as A
import matplotlib.pyplot as plt
import matplotlib.patches as patches
image = plt.imread('sample.jpg')
xmin, ymin, xmax, ymax = [165, 264, 253, 372]
bbox = [[xmin, ymin, xmax, ymax, 'chair']]
transform = A.Compose([
A.HorizontalFlip(p=1),
], bbox_params=A.BboxParams(format='pascal_voc'))
transformed = transform(image=image, bboxes=bbox)
transformed_image = transformed['image']
transformed_bboxes = transformed['bboxes']
# 이미지 출력
plt.imshow(transformed_image)
# bounindg box 그리기
for transformed_bbox in transformed_bboxes:
xmin, ymin, xmax, ymax, category = transformed_bbox
rect = patches.Rectangle((xmin, ymin), xmax-xmin, ymax-ymin, linewidth=2, edgecolor='red', fill=False)
ax = plt.gca()
ax.add_patch(rect)
plt.xticks([]); plt.yticks([])
plt.show()
