Object Detection 정리 3. Grid for Muliple Objects 및 Non-Maximum Suppresion (NMS)

AI 딥러닝을 통한 Object Detection에 대해서 간단하게 정리해보자.

지난 정리 2에서는 Localization 모델에 Confidence Score를 추가하여 오브젝트가 없는 경우에 대응할 수 있도록 했다.

문제. 오브젝트가 여러개라면?

 예제와 같이 오브젝트가 2개인 이미지를 우리 모델은 제대로 예측해낼 수 없다. 오직 BBOX 하나에 대한 위치정보만 return하기 때문이다. 해서 아마도 우리 모델은 두 오브젝트의 중간값 혹은 평균값을 return할 것이다.

오브젝트가 2개인 이미지(left)에 대한 예측은 아마도 두 오브젝트의 중간을 return할 것(right)이다.

해결. 그리드 방식으로 접근하자.

 그리드 방식은 이미지를 그리드로 나눠 각 그리드에서 각각 위치 정보를 return한다.

예.) 총 16개의 그리드로 나눠 16개의 위치정보를 return한다. (NMS output 추가??)

 상기 예제 이미지처럼 4x4그리드로 나눠서 위치 정보를 return하기 위해서 모델의 Head 구조를 수정해야하며, 또 데이터에 라벨링된 위치정보도 동 그리드 구조에 맞춰 전처리해줘야 한다. 즉, 모델의 Output Dimension이 (None, 4, 4, 5)여야 한다. 그리고 모델의 학습을 위해서는 총 16개의 그리드 중 실제 오브젝트(Ground Truth, 줄여서 GT라고 한다)의 중심점이 있는 그리드에서만 Loss를 계산하고 그렇지 않은 그리드는 무시해야 한다.

# 모델 구조
import tensorflow as tf

backbone = tf.keras.applications.mobilenet.MobileNet(
    include_top=False,
    weights=None,
    pooling=None,
)

neck = tf.keras.layers.MaxPooling2D()
x = neck(backbone.output)

head = tf.keras.layers.Conv2D(
    filters=5,
    kernel_size=3,
    padding='same',
    activation='sigmoid',
)

model = tf.keras.Model(
    inputs=backbone.input,
    outputs=head(x)
)
model.summary()

# 데이터 라벨 전처리
import cv2
import numpy as np

output_shape = (4,4,5)
x_grid, y_grid = 1/output_shape[1], 1/output_shape[0]
label = np.zeros(output_shape, dtype=np.float32)
for box in boxes:
    x, y, w, h = box
    x_cell, y_cell = int(x//x_grid), int(y//y_grid)
    label[y_cell, x_cell, 0] = 1.0
    label[y_cell, x_cell, 1] = x/x_grid - x_cell
    label[y_cell, x_cell, 2] = y/y_grid - y_cell
    label[y_cell, x_cell, 3] = w
    label[y_cell, x_cell, 4] = h

# Loss Function과 학습
def grid_loss(y_true, y_pred):
    true_conf = y_true[..., 0]
    loss = tf.math.reduce_mean(tf.math.abs(y_pred - y_true), axis=-1)
    loss = tf.where(tf.equal(true_conf, 1.0), loss, 0.0)
    return tf.math.reduce_mean(loss)

 이렇게 그리드로 나눠 학습하고 예측하니 같은 오브젝트에 대해서 여러 개의 위치 정보가 나오게 되는데, 우리가 의도하는 기능은 한 오브젝트에 하나의 바운드 박스를 그리는 것이다. 따라서 겹치는 여러 개의 Box 중 가장 높은 Confidence Score를 가지는 Box만 남기고 무시하도록 하는 알고리즘인 Non-Maximum Suppresion(NMS)을 사용한다.

# NMS
def nms(pred, thresh=0.5):
    xy1 = pred[..., 1:3] - pred[..., 3:]/2
    xy2 = pred[..., 1:3] + pred[..., 3:]/2
    bbox = np.concatenate([pred[..., 0:1], xy1, xy2], axis=-1)
    bbox = bbox[pred[...,0]>thresh]
    bbox = np.sort(bbox, axis=0)[::-1]
    if bbox.shape[0] == 0:
    	return []
    
    mask = []
    conf, x1, y1, x2, y2 = bbox[:, 0], bbox[..., 1], bbox[..., 2], bbox[..., 3], bbox[..., 4]
    area = (x2-x1) * (y2-y1)
    idxs = np.argsort(conf)
    while len(idxs)>0:
    	last = len(idxs) - 1
        i = idxs[last]
        mask.append(i)
        
        xx1 = np.maximum(x1[i], x1[idxs[:last]])
        yy1 = np.maximum(y1[i], y1[idxs[:last]])
        xx2 = np.minimum(x2[i], x2[idxs[:last]])
        yy2 = np.minimum(y2[i], y2[idxs[:last]])
        w, h = np.maximum(0, xx2-xx1+1), np.maximum(0, yy2-yy1+1)
        iou = (w*h) / area[idxs[:last]]
        idxs = np.delete(idxs, np.concatenate(([last], np.where(iou > thresh)[0])))
	return bbox[mask]

NMS를 이용하여 겹치는 BBOX 중 Confidence Score가 가장 높은 BBOX만 남긴다.