Albumentation을 이용해 Image augmentation 해보기03- 바운딩박스 포함 증강

이번 글에서는 albumentations.ai/docs/examples/example_bboxes/ 에 있는

object detection용 바운딩박스 증강 예제를 연습해본다.

 

1. 먼저 필요한 라이브러리들을 import 해준다.

import cv2
import albumentations as A
import matplotlib.pyplot as plt
import numpy as np
import random

---

2. 이미지에서 경계상자 및 클래스 레이블을 시각화 하는 함수를 정의한다

OX_COLOR = (255, 0, 0) # Red
TEXT_COLOR = (255, 255, 255) # White

def visualize_bbox(img, bbox, class_name, color=BOX_COLOR, thickness=2):
    """Visualizes a single bounding box on the image"""
    x_min, y_min, w, h = bbox
    x_min, x_max, y_min, y_max = int(x_min), int(x_min + w), int(y_min), int(y_min + h)

    cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color=color, thickness=thickness) 
	
    #이부분 이해하기
    ((text_width, text_height), _) = cv2.getTextSize(class_name, cv2.FONT_HERSHEY_SIMPLEX, 0.35, 1)
    
    cv2.rectangle(img, (x_min, y_min - int(1.3 * text_height)), (x_min + text_width, y_min), BOX_COLOR, -1)
    cv2.putText(
        img,
        text=class_name,
        org=(x_min, y_min - int(0.3 * text_height)),
        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
        fontScale=0.35, 
        color=TEXT_COLOR, 
        lineType=cv2.LINE_AA,
    )
    return img
 
def visualize(image, bboxes, category_ids, category_id_to_name):
    img = image.copy()
    for bbox, category_id in zip(bboxes, category_ids):
        class_name = category_id_to_name[category_id]
        img = visualize_bbox(img, bbox, class_name)
    plt.figure(figsize=(12, 12))
    plt.axis('off')
    plt.imshow(img)
    plt.show()

visualize_bbox:

- 함수의 인자로 이미지,  바운딩박스(1차원배열), 클래스명(문자열), 색상(튜플)을 받는다.

- 2차원 배열 내  원소마다 [x_min, y_min, x_max, y_max]가 담겨있다. 이를 통해서 바운딩박스의 시작과 끝점 좌표를정의한다.

- cv2.getTextSize(텍스트, 폰트, 폰트 크기, 두께)를 입력해서  폰트size(x,y) 튜플,

베이스라인(맨아래 텍스트 지점을 기준으로 한 기준선의 y좌표)을 ( (w,h), _)로 리턴받는다.  

 

- cv2.putText로( 클래스명, 문자열을 출력할 위치의 좌측 하단 좌표(x,y)좌표, 폰트 종류, 폰트크기 ,색, 두께, 라인 타입)을 받아 작성한다.

 

# 이미지 및 주석 얻기
image = cv2.imread('coco.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
bboxes = [[5.66, 138.95, 147.09, 164.88], [366.7, 80.84, 132.8, 181.84]]
category_ids = [17, 18]
category_id_to_name = {17: 'cat', 18: 'dog'}

visualize:

- 함수의 인자로 이미지, 바운딩박스(2차원배열),  클래스아이디(1차원배열) , 클래스아이디(딕셔너리) 를 받는다

- bboxes(2차원배열), category_dis(1차원배열)의 각 원소마다

class name는  category_id_to_name의 category_ids의 원소에 해당하는 키값을 가진 value의 문자열이 되고

이미지에는 각 바운딩박스에 각 클래스명에 해당하는 레이블과 바운딩박스를 그린다.

- plt로 영상을 출력한다.

 

---

3. 정의한 변수들과 함수들을 가지고 바운딩박스 처리된 이미지를 출력한다.

visualize(image, bboxes, category_ids, category_id_to_name)

이렇게 바운딩 박스가 처리된 영상이 출력된다.

---

 

4. Albumentations로 증강 파이프라인을 작성한다

#증강 파이프라인 정의
transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.RandomContrast(limit=0.5,p=1)
    ],
    bbox_params=A.BboxParams(format='coco', label_fields=['category_ids']),
)
#시각화 목적으로 랜덤 시드를 수정해서 증강된 화면이 일정하게 나오도록 설정
random.seed(7)
transformed = transform(image=image, bboxes=bboxes, category_ids=category_ids)
visualize(
    transformed['image'],
    transformed['bboxes'],
    transformed['category_ids'],
    category_id_to_name,
)

- Compose안에 배열로 여러 증강 방법들을 넣는건 해봤지만 그 외의 매개변수가 들어가는것을 처음봐서  API를 다시 뜯어봤다. 

- bbox_params(format=" 'coco', 'pascal_voc', 'albumentations'또는 'yolo' ", label_field=[바운딩박스의 라벨링 키값 배열]

을 추가해주므로써 바운딩박스를 함께 증강시켜준다.

 

- transformed는 transform(Compose클래스 객체를 ( image = coco사진, bbox= bboex배열, 라벨 키= cotegort_ids)로 생성 한 후에 

- visualize() 함수에 compose객체의 (이미지, 바운딩박스 , 라벨배열, 라벨 대응텍스트 딕셔너리)를 넣어줘서 이미지를 출력한다

 

---

4. 증강절차를 거친 이미지다.  HorizontalFlip과 RandomContrast가 잘 적용된 모습이고 바운딩박스가 유지되었다.