[Python] Object Detection Mosaic Augmentation :: YOLO v5

Mosaic 이란?

Mosaic 기법은 Image Augmentation 기법 중 하나로, 4장의 이미지를 한 장으로 만드는 기법이다.
예를 들면, 다음 4장의 이미지를

다음과 같이 한 장으로 만든다.

각 이미지의 사이즈는 random하게 하여 자유도를 준다.
따라서 전체 이미지가 모두 포함되지 않을 수도 있다.
(빨간색 박스는 Ground Truth이다. 알아보기 좋게 같이 그렸다.)

이 기법은 이미지 사이즈가 작아져, 작은 물체가 많아지므로 작은 물체를 Detect하기 어렵다는 YOLO의 단점을 극복 할 수 있도록 도와주며, 1장의 이미지로 4장의 이미지를 학습하는 효과를 주어 적은 batch size로도 학습이 용이하도록 해준다.

---

YOLOv5에 저자가 직접 구현해 놓은 Mosaic 기법을 조금 변형하여 Offline Mosaic Augmentation 코드를 짜보았다.
YOLOv5에서는 Online Mosaic으로, 학습시킬 때 자동으로 현재 batch에서 다루고있는 이미지 + random 3장 이미지를 뽑아서 Mosaic 이미지를 만든다.

변경한 코드는 4장 모두 random하게 뽑아서 이미지를 만든다.

코드는 YOLOv5 utils dataloaders.py에 load_mosaic()으로 구현되어있다.
이는 학습 할 때 train.py 코드에 train_loader, dataset 만들때 사용된다.

참고로 아래 이미지에서 line 228에 augment=True로 되어있다.
이는 학습 시킬 때 따로 파라미터로 입력하지는 않는 부분이라, Online Augment를 원치 않는다면 이부분을 False로 직접 바꿔주어야 한다.

본인은 데이터 특성 상 Mosaic을 원치 않는 Class가 있어, Online Mosaic을 막고 Offline Mosaic으로 수행하기 위해 코드를 수정했다.

YOLO v5 저자가 작성해놓은 코드를 거의 그대로 사용하되, 조금 변형했다.
본 포스팅에서 그 코드를 소개하고 설명하고자 한다.

---

최종 코드는 다음과 같다.

import math
import os
import random
from pathlib import Path
import numpy as np
import yaml
import glob
import cv2
import torch
from pathlib import Path
import pandas as pd

def img2label_paths(img_paths):
    # Define label paths as a function of image paths
    sa, sb = f'{os.sep}images{os.sep}', f'{os.sep}labels{os.sep}'  # /images/, /labels/ substrings
    return [sb.join(x.rsplit(sa, 1)).rsplit('.', 1)[0] + '.txt' for x in img_paths]
def myLoadImage(i, im_files, augment=True) :
    f = im_files[i]
    im = cv2.imread(f)
    h0, w0 = im.shape[:2]
    r = img_size / max(h0, w0)
    if r != 1:  # if sizes are not equal
        interp = cv2.INTER_LINEAR if augment else cv2.INTER_AREA  # random.choice(self.rand_interp_methods)
        im = cv2.resize(im, (int(w0 * r), int(h0 * r)), interpolation=interp)
    return im, (h0, w0), im.shape[:2] 
def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
    # Convert nx4 boxes from [x, y, w, h] normalized to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = w * (x[:, 0] - x[:, 2] / 2) + padw  # top left x
    y[:, 1] = h * (x[:, 1] - x[:, 3] / 2) + padh  # top left y
    y[:, 2] = w * (x[:, 0] + x[:, 2] / 2) + padw  # bottom right x
    y[:, 3] = h * (x[:, 1] + x[:, 3] / 2) + padh  # bottom right y
    return y
def xyxy2xywhn(x, w=640, h=640, clip=False, eps=0.0):
    # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] normalized where xy1=top-left, xy2=bottom-right
    if clip:
        clip_coords(x, (h - eps, w - eps))  # warning: inplace clip
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = ((x[:, 0] + x[:, 2]) / 2) / w  # x center
    y[:, 1] = ((x[:, 1] + x[:, 3]) / 2) / h  # y center
    y[:, 2] = (x[:, 2] - x[:, 0]) / w  # width
    y[:, 3] = (x[:, 3] - x[:, 1]) / h  # height
    return y
def noPad(label, xc, yc, xmin, ymin) :
    label4_ = label.copy()
    if ((xc>s) & (yc<s)) :
        label4_[:,1] = label4_[:,1] - xmin
        label4_[:,3] = label4_[:,3] - xmin
    elif ((xc<s) & (yc>s)) :
        label4_[:,2] = label4_[:,2] - ymin
        label4_[:,4] = label4_[:,4] - ymin
    elif ((xc>s) & (yc>s)) :
        label4_[:,1] = label4_[:,1] - xmin
        label4_[:,3] = label4_[:,3] - xmin
        label4_[:,2] = label4_[:,2] - ymin
        label4_[:,4] = label4_[:,4] - ymin
    return label4_

def myMosaic(img_size, path, form = "xyxy") :
    labels4 = []
    s = img_size
    mosaic_border = [-img_size // 2, -img_size // 2]

    # center point
    yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in mosaic_border)  # mosaic center x, y # 81 172

    # image files and labels
    # train set & not Scratch
    im_files = np.array([i for i in os.listdir(path) if i.endswith("jpg")], dtype = "object")
    im_files = list(path + "/" + im_files)

    label_files = img2label_paths(im_files) 

    n = len(im_files)
    indices_ = range(n)

    # random 4 images
    indices = random.choices(indices_, k=4)  # 3 additional image indices # [0, 3292, 20762, 18713]
    random.shuffle(indices)  # [18713, 0, 20762, 3292]

    for i, index in enumerate(indices):
        # Load image
        img, (h0, w0), (h, w) = myLoadImage(index,im_files)

        # place img in img4
        if i == 0:  # top left
            img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)  # base image with 4 tiles
            x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)
            x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (small image)

            xmin, ymin = x1a, y1a

        elif i == 1:  # top right
            x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
            x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
        elif i == 2:  # bottom left
            x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
            x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
        elif i == 3:  # bottom right
            x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
            x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)

            xmax, ymax = x2a, y2a

        img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b]  # img4[ymin:ymax, xmin:xmax]
        padw = x1a - x1b
        padh = y1a - y1b

        # Labels
        labels = np.loadtxt(label_files[index], ndmin=2)

        if labels.size:
            labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh)  # normalized xywh to pixel xyxy format
            labels[labels < 0] = 0
            labels[labels > s*2] = s*2
        labels4.append(labels)

    # Concat/clip labels
    labels4 = np.concatenate(labels4, 0)

    # no padding
    img4 = img4[ymin:ymax, xmin:xmax,:]
    labels4 = noPad(labels4, xc, yc, xmin, ymin)

    # resize
    img4 = cv2.resize(img4, (w,h))
    labels4[:,1:] = xyxy2xywhn(labels4[:,1:], w = xmax - xmin, h = ymax-ymin)
    labels4 = labels4[labels4[:,3]>0]
    labels4 = labels4[labels4[:,4]>0]
    
    if form == "xyxy" :
        labels4[:,1:] = xywhn2xyxy(labels4[:,1:], w = img4.shape[1], h = img4.shape[0])
    
    return img4, labels4

from PIL import Image, ImageDraw
import matplotlib.patches as patches
def myFig(img, bbox, drawB = True) :
    image = Image.fromarray(img,"RGB")
    draw = ImageDraw.Draw(image)

    if drawB :
        for i in range(len(bbox)) :
            xmin = bbox[i][1] 
            ymin = bbox[i][2]
            xmax = bbox[i][3]
            ymax = bbox[i][4]

            draw.rectangle((xmin, ymin, xmax, ymax), outline=(255,0,0), width = 1) # bounding box
#     display(image)
    return image

마지막에 myFig function은 이미지를 그려보기 위해 만든 function이다.

실행한 결과는 다음과 같다.

myMosaic function은 4장의 이미지를 합친 1장의 image와 label(class, GT)를 return한다.
위 이미지는 image와 label을 같이 그린 결과이다.

(Open cv 컬러 이미지를 PIL로 그릴 땐 cvtColor로 채널 순서를 바꿔주어야한다.)

---

이제 코드를 자세히 보겠다.
먼저 img_size와 path를 setting해준다.

img_size = 320
path = "YOUR PATH"

img_size는 최종 Mosaic된 이미지의 Width와 Height 중 max 값을 img_size로 지정하겠다는 뜻이다.
path는 jpg 이미지와 label txt 파일이 있는 경로를 의미한다.
coco128 train 2017 데이터를 사용하겠다.

참고로 해당 path에 데이터는 다음과 같이 이미지와 label txt파일이 같이 들어있어야 한다.

label 형태는 기본 YOLO 형태와 같다.

이제 최종 4장 이미지의 label(class, GT)을 담을 변수와 기타 필요한 변수들을 선언한다.

labels4 = [] 
s = img_size
mosaic_border = [-img_size // 2, -img_size // 2]

다음으로 image와 label 파일 목록을 불러온다.

im_files = np.array([i for i in os.listdir(path) if i.endswith("jpg")], dtype = "object")
im_files = list(path + "/" + im_files)

label_files = img2label_paths(im_files)

img2label_paths()는 위에서 미리 선언한 function을 사용한다.

다음으로 파일 개수를 range로 선언해, random 한 4장 이미지를 뽑을 준비를 한다.

n = len(im_files)
indices_ = range(n)

4장 이미지를 합칠 center point 좌표를 random하게 뽑는다.

# center point
# random.seed(1)
yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in mosaic_border)

(166, 462) 위치를 기준으로 왼쪽 위에 첫 번째 이미지, 오른쪽 위에 두 번째 이미지, 아래 왼쪽에 세 번째 이미지, 아래 오른쪽에 네 번째 이미지를 삽입한다는 의미이다.

최종 이미지 사이즈는 320이지만, 이미지를 합칠 때는 320*2 by 320*2 사이즈에 합치고 다음에 resize하기 때문에 center point가 320을 넘어갈 수도 있다.

삽입할 4장의 이미지를 random하게 뽑고 순서를 섞는다.

# random 4 images
indices = random.choices(indices_, k=4) 
random.shuffle(indices)

이제 이 4장의 이미지를 위에서 미리 뽑은 center point를 기점으로 합친다.

for i, index in enumerate(indices):
    # Load image
    img, (h0, w0), (h, w) = myLoadImage(index,im_files)

    # place img in img4
    if i == 0:  # top left
        img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)  # base image with 4 tiles
        x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)
        x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (small image)

        xmin, ymin = x1a, y1a

    elif i == 1:  # top right
        x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
        x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
    elif i == 2:  # bottom left
        x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
        x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
    elif i == 3:  # bottom right
        x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
        x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)

        xmax, ymax = x2a, y2a

    img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b]  # img4[ymin:ymax, xmin:xmax]
    padw = x1a - x1b
    padh = y1a - y1b

    # Labels
    labels = np.loadtxt(label_files[index], ndmin=2)

    if labels.size:
        labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh)  # normalized xywh to pixel xyxy format
        labels[labels < 0] = 0
        labels[labels > s*2] = s*2
    labels4.append(labels)

---

이를 i=0부터 하나하나 보면 다음과 같다.

먼저, i=0일 때 img4라는 변수를 선언한다. 최종적으로 widgh, height 중 max가 320인 이미지를 만들거고,
320 size 이미지를 4장 붙이기 위해 320*2 by 320*2 size의 빈 array를 만든다. 값은 114로 채웠으며, 이를 시각화하면 회색이다.

그리고 그 img4에 위에서 미리 뽑은 center point를 기준으로 왼쪽 위에 첫 번째 이미지를 삽입한다.

참고로 원래 이미지는 다음과 같다.

i=1일 때 두 번째 이미지를 img4에 center point를 기준으로 오른쪽 위에 삽입한다.

i=2일 때 왼쪽 아래에 삽입한다.

마지막 이미지는 왼쪽 아래 삽입한다.

---

이제 list로 append했던 4장 이미지에 대한 label을 numpy로 합쳐주고,

# Concat/clip labels
labels4 = np.concatenate(labels4, 0)

이미지에 왼쪽과 아래 빈 부분을 잘라준다.

# no padding
img4 = img4[ymin:ymax, xmin:xmax,:]
labels4 = noPad(labels4, xc, yc, xmin, ymin)

원래 YOLO v5 코드에 이 부분은 없고, 대신 random 하게 crop하고 기존에 정해준 image size대로 resize해주는 부분이 있다.
본인은 데이터 특성 상 이미지가 너무 많이 잘리는걸 원치 않아서 padding 된 부분만 잘라주고, resize하는 것으로 코드를 수정했다.

마지막으로 처음에 입력해준 image size로 resize해주고,
label도 같이 변경해준다.

img4 = cv2.resize(img4, (w,h))
labels4[:,1:] = xyxy2xywhn(labels4[:,1:], w = xmax - xmin, h = ymax-ymin)
labels4 = labels4[labels4[:,3]>0]
labels4 = labels4[labels4[:,4]>0]
labels4[:,1:] = xywhn2xyxy(labels4[:,1:], w = img4.shape[1], h = img4.shape[0])

최종적으로 return된 이미지의 height가 320이고, Ground Truth도 맞춰서 잘 변경된 것을 볼 수 있다.

Mosaic_Reference.ipynb

2.36MB