Object Detection 정리 4. Classification을 Localization에 추가

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

지난 정리 3에서는 Grid 방식을 적용하여 오브젝트가 다수인 경우에 대응할 수 있도록 했다.

문제. 그래서 찾은 오브젝트가 무엇인가?

 우리 모델의 예측은 아직 위치 정보만 포함하고 있다.(Localization Only) 즉, 찾아낸 오브젝트가 무엇인지 분류하지 않는다. 이번 포스트에서는 우리 모델이 위치 정보와 함께 분류 정보도 리턴하도록 수정한다.(Detection = Localization + Classification)

해결. 모델 Output Dimension에 Classification Dimension을 추가하자.

 분류 모델은 클래스 개수만큼의 Output Dimension을 가진다. 클래스가 [남자, 여자, 기타] 3가지라면 Output Dimension은 3이어야 한다. 이 3을 우리 모델의 Output Dimension에 추가하여, 같은 Feature Map에서 Localization과 동시에 Classification을 수행할 것이다. Localization과 Classification을 동시에 수행하는 아웃풋은 [conf, x, y, w, h, 남(Male), 여(Female), 기타(Etc.)]의 형태를 가진다. 정리하자면 conf가 threshold 이상이면 x,y,w,h를 이용해 바운드 박스를 그리고 male,female,etc.를 비교해 가장 값이 높은 클래스로 박스에 네이밍을 한다.

Output [conf, x, y, w, h, male, female, etc.] 중 첫 5개 [conf, x, y, w, h]로 Bbox를 그리고, 뒤 3개 [male, female, etc.]로 Classification을 수행한다.

 이전 포스트에서와 같이 오브젝트가 다수인 경우에 대응하기 위해서 그리드 방식으로 Output Dimension을 확장해서 Object Detection을 최종 수행할 수 있게 됩니다.

예제 이미지에서 최종 Output Dimension은 (4,4,8)이다. 각각 행, 열, 정보를 뜻한다.

# 모델 구조
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=8,
    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,8)
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
    if cls = 'male':
    	label[y_cell, x_cell, 5] = 1.0
    elif cls == 'female':
    	label[y_cell, x_cell, 6] = 1.0
    else:
    	label[y_cell, x_cell, 7] = 1.0

 이제 Object Detection 모델을 완성했습니다. 꽤 괜찮은 성능을 보여줍니다. 다음부터는 논문을 참조하여 성능을 향상시키는 기법들을 실험해보겠습니다.