YOLOv5 이미지 데이터셋을 구축해보자! - Imgaug 사용법 (Bounding Box 같이 회전시키기)

이전 글에서는 Bounding Box를 그리는 labelImg 라이브러리에 대해 알아보았습니다. 이번 글에서는 Bounding Box를 그린 이미지에 대해 Augmentation을 수행해보겠습니다. 

만약 Bounding Box를 직접 그리는 방법이 궁금하시면 아래 글을 참고해주세요!

https://aiclaudev.tistory.com/29

YOLOv5 이미지 데이터셋을 구축해보자! - labelImg 사용법

 

 

https://imgaug.readthedocs.io/en/latest/index.html

imgaug — imgaug 0.4.0 documentation

위 링크는 imgaug documentation입니다. install이나 다양한 augmentation parameter 등이 나와있으니 참고하시는걸 적극 추천드리고, 위 documentation에서 굉장히 많은 부분을 참고하여 작성하였음을 말씀드립니다.

 

 

저같은 경우 Jupyter notebook 환경에서 Imgaug를 사용했습니다. 셀을 실행할때마다 결과를 확인할 수 있어서, 다양한 Augmentation 인자들을 입력해가면서 결과를 확인하기 확실히 편합니다!

Anaconda prompt 등에서 아래와 같은 명령어를 실행하여 Imgaug 라이브러리를 우선 install 해줍니다.

pip install imgaug

 

귀찮으시다면 그냥 Jupyter notebook 상에 있는 셀에서 아래와 같은 명령어를 실행해주셔도 무방합니다.

!pip install imgaug

 

다음으로, Augmentation을 수행하기 위해 몇가지 함수를 정의하겠습니다. 각각 readBound, convert, revert입니다. 아래 이미지는 캡쳐되어있는 이미지이고, 복사 가능한 코드는 맨 아래쪽에 모아둘테니 참고해주세요!

 

① readBound

함수명 그대로, 이미지의 BoundingBox좌표를 가져오는 함수이고 코드는 아래와 같습니다. 

txt파일 내에 Class number로 같이 있는 걸 기억하실 거에요. 주의하실 점은 함수의 return은 Class number까지 포함하고 있고(즉 Bounding Box의 좌표만을 return하지 않는다는 의미입니다.) 이에 대한 처리는 convert와 revert에서 해줍니다.

 

② convert, revert

위 함수가 하는 역할에 대해서 먼저 설명하자면, YOLO의 Bounding Box 좌표를 변환해주는 역할입니다.

이게 무슨 뜻이냐 하면, YOLO의 Bounding Box좌표는 x_start, y_start, width, height로 정의되어 있고 이미지 전체에 대한 상대적인 좌표(?)로 정의됩니다.

하지만, imgaug 라이브러리에서 Bounding Box의 좌표는 x_start, y_start, x_end, y_end와 같이 구성되어있고, 이미지 전체에 대한 상대적인 비율이 아닌 이미지 크기에 따른 절대적인 좌표(?)로 표현됩니다.

 

예를 들어서 100*100 이미지가 있다고 해봅시다.  Bounding Box의 왼쪽 아래 꼭짓점이 전체 이미지의 1/4에 위치하고, 오른쪽 위에 꼭짓점이 전체 이미지의 3/4에 위치한다고 하겠습니다.

이때 YOLOv5의 Bounding Box는 (0.25, 0.25, 0.5, 0.5)로 표현됩니다. 즉, 각각 숫자가 x_start, y_start, width, height를 의미하고, 전체 이미지에 상대적인 크기로 기록된다는 것이죠.

하지만 imgaug에서 Bounding Box는 (25, 25, 75, 75)로 표현됩니다. 즉, 각각 숫자가 x_start, y_start, x_end, y_end로 표현되고 YOLOv5와 달리 이미지 크기에 대한 좌표로 표현됩니다.

 

따라서, 위와 같은 처리를 위해 아래와 같은 두 함수를 정의해주어야 합니다!

convert는 YOLOv5의 상대적인 좌표를 imgaug에서 사용할 수 있도록 절대좌표로 변환하는 역할을 수행합니다. parameter로 size와 box가 있는 것을 보실 수 있네요. size란 이미지의 크기를 의미하고 box는 이미지의 Bounding Box 상대적인 좌표를 의미합니다. 즉, 위에서 다룬 readBound의 return값이라고 생각하시면 됩니다! imgaug는 절대좌표를 사용하기에 size라는 이미지 크기를 곱해주는 것을 보실 수 있습니다.

 

revert는 imgaug에서의 Bounding Box를 YOLOv5의 상대적인 좌표로 다시 변환하는 역할을 합니다. parameter로 size와 box를 보실 수 있습니다. size는 이전과 같이 이미지의 크기를 의미하며, box는 Bounding Box의 좌표를 의미합니다. convert의 box와는 조금 다릅니다. convert의 box는 'YOLOv5가 사용하는 상대적인 좌표'라고 생각하시면 되고, 이번 revert의 box는 'imgaug가 사용하는 절대적인 좌표'라고 생각하시면 될 것 같습니다. 

 

 

정리하자면, 위 세가지 함수를 사용하여 Augmentation을 사용하는 과정을 아래와 같습니다.

(1) readBound로 이미지의 Bounding Box의 상대적인 좌표(x_start, y_start, width, height) 추출하기

(2) convert로 Bounding Box의 상대적인 좌표를 imgaug에서 사용할 수 있는 절대적인 좌표(x_start, y_start, x_end, y_end)로 바꾸기

(3) Augmentation을 한 후, 절대적인 좌표 상태인 Bounding Box를 다시 YOLOv5에서 사용할 수 있게 상대 좌표로 변환하기.

 

 

이제 거의 다 왔습니다. 사실 imgaug라이브러리를 사용하는 것은 전혀 어렵지가 않아요. 단지 이미지 증강 시에 Rotation이나 확대/축소 등을 할 때면 Bounding Box의 좌표 또한 달라져야 하기 때문에 위 함수들을 정의했을 뿐입니다.

 

 

이제 실제 Augmentation을 진행하는 코드 부분을 보도록 하겠습니다. 다양한 인자들이 iaa.Sequential에 정의되어 있는 것을 보실 수 있네요. Multiply, LaplaceNoise 등등... 사실 LaplaceNoise와 같은 인자들은 Bounding Box의 좌표를 변환시키지는 않습니다. 하지만, 이미지를 확대/축소하거나 회전하는 경우 Bounding Box의 좌표가 당연히 바뀌어야하겠죠. 위에 정의되어있는 인자들을 제외한 다른 인자들은 Documentation에서 참고하도록 해주세요!

https://imgaug.readthedocs.io/en/latest/index.html

imgaug — imgaug 0.4.0 documentation

 

 

 

 

이제 모든 준비가 완료되었으니, 전체 코드를 보면서 활용해보도록 합시다. 필요하시다면 언제든지 복붙해서 사용하셔도 괜찮습니다.

def readBound(imgname) :
    with open("파일경로 and 파일명.txt", 'r') as f :
        s = f.read()
    s = s.replace('\n', '')
    result = s.split(' ')
    for i in range(len(result)) :
        result[i] = float(result[i])
    return result

 

def convert(size, box) :
    x1 = (box[1] - box[3] / 2) * size[0]
    x2 = (box[1] + box[3] / 2) * size[0]
    y1 = (box[2] - box[4] / 2) * size[1]
    y2 = (box[2] + box[4] / 2) * size[1]
    return [x1, y1, x2, y2]

 

def revert(size, box): # box : [x1,y1,x2,y2]
    dw = 1./size[0] 
    dh = 1./size[1] 
    x = (box[0] + box[2])/2.0 
    y = (box[1] + box[3])/2.0 
    w = box[2] - box[0] 
    h = box[3] - box[1] 
    x = x*dw 
    w = w*dw 
    y = y*dh 
    h = h*dh 
    return [x,y,w,h]

 

import imageio
import imgaug as ia
import imgaug.augmenters as iaa
from imgaug.augmentables.bbs import BoundingBox, BoundingBoxesOnImage
%matplotlib inline

filepath = "이미지가 있는 폴더 경로"

def augimg(idx, imgname) :
        bound = readBound(imgname) # 상대좌표 (x_start, y_start, width, height)를 읽어옵니다.
        class_num = int(bound[0]) # class_num에 클래스를 저장합니다.
        image = imageio.imread(filepath+"/"+imgname+".jpg") # 이미지를 불러옵니다.
        size = image.shape # 불러온 이미지의 size를 저장합니다.
        point = convert(size, bound) # 상대좌표 (x_start, y_start, width, height) -> 절대좌표 (x_start, y_start, x_end, y_end)
        
        bbs = BoundingBoxesOnImage([
            BoundingBox(x1=point[0], y1=point[1], x2=point[2], y2=point[3])
        ], shape=image.shape)

        # 실제 Augmentation 인자를 추가하는 부분입니다. Documentation을 참고하여 원하는 인자를 선택해주세요!
        seq = iaa.Sequential([
            iaa.Multiply((0.8, 1.1)),
            iaa.AdditiveLaplaceNoise(scale=(10,30), per_channel=True),
            iaa.AdditiveGaussianNoise(scale=(10, 80)),
            iaa.Crop(percent=(0, 0.2)),
            iaa.Affine(
                translate_px={"x": 40, "y": 60},
                scale=(0.5, 0.7),
                rotate=(-70, 70))
        ])

        # 이미지에 대해 Augmentation을 진행합니다. Rotation과 같은 인자가 있다면 Bounding Box의 좌표 역시 바뀌어야 합니다.
        image_aug, bbs_aug = seq(image=image, bounding_boxes=bbs)
       
        before = bbs.bounding_boxes[0]
        after = bbs_aug.bounding_boxes[0]
        
        aug_box = [after.x1, after.y1, after.x2, after.y2]
        aug_box = revert(size, aug_box)
        aug_box = [class_num] + aug_box
        
        for i in range(1, len(aug_box)) :
            aug_box[i] = round(aug_box[i], 6)
        for i in range(len(aug_box)) :
            aug_box[i] = str(aug_box[i])
        
        # 증강된 이미지에 대한 상대적인 Bounding Box좌표(x_start, y_start, width, height)를 txt파일 형태로 저장합니다. 
        with open(f"{filepath}/{imgname}{idx}"+".txt", 'w') as f :
            f.writelines([aug_box[0], ' ', aug_box[1], ' ', aug_box[2] , ' ', aug_box[3], ' ', aug_box[4], '\n'])
        
        # 증강된 이미지를 저장합니다.
        imageio.imwrite(f"{filepath}/{imgname}{idx}"+".jpg", image_aug)

 

imgs = ['apple', 'banana', 'tomato'] # file path 내에 있는 이미지의 이름입니다.
for i in imgs :
    for j in range(500) : # file path내에 있는 각 이미지를 500장씩 증강합니다.
        augimg(j, i)

 

 

 

 

잘못된 내용에 대한 피드백은 언제나 환영합니다!