YOLO 모델로 객체 검출 시작하기, YOLO 출력 후처리 비최대 억제(NMS) 기법 설명과 적용 방법

YOLO(You Only Look Once)는 객체 검출 분야에서 매우 인기 있는 알고리즘 중 하나입니다. YOLO는 이미지를 단일 네트워크에서 처리하여 객체의 위치와 종류를 동시에 예측하는 혁신적인 접근 방식을 제공합니다. 기존의 객체 검출 알고리즘들이 이미지를 여러 부분으로 나눈 후 각각의 부분에 대해 클래스를 예측하던 방식과는 달리, YOLO는 전체 이미지를 통해 객체를 인식합니다. 이 장에서는 YOLO의 기본 개념과 함께 YOLO의 출력 후처리 과정인 비최대 억제(NMS)에 대해 살펴보겠습니다.

YOLO의 기본 구조

YOLO는 CNN(Convolutional Neural Network)을 기반으로 하며, 입력 이미지를 그리드로 나누고 각 그리드 셀에 대해 객체 존재 가능성과 바운딩 박스 정보를 추정합니다. YOLO는 다음의 주요 단계를 통해 작동합니다:

  1. 입력 이미지가 네트워크를 통과하며 피쳐를 추출합니다.
  2. 네트워크의 마지막 레이어에서 그리드 셀 수 만큼의 바운딩 박스와 확률을 출력합니다.
  3. NMS를 사용하여 중복된 박스를 제거하고 최종 예측 결과를 도출합니다.

YOLO 출력 해석

YOLO는 각 바운딩 박스에 대해 다음과 같은 정보를 출력합니다:

  • 박스의 위치 (X, Y, Width, Height)
  • 객체의 클래스 가능성 (Probability)
  • 각 클래스에 대한 확률 스코어

이러한 출력 데이터는 최종적으로 후처리 단계를 거쳐 정리됩니다.

비최대 억제(NMS)란?

비최대 억제(NMS)는 여러 개의 중복된 바운딩 박스를 정리하기 위해 사용되는 후처리 기법입니다. YOLO의 경우, 동일한 객체에 대해 여러 개의 바운딩 박스가 있을 수 있는데, 이는 같은 객체를 여러 번 검출하는 경우에 해당합니다. NMS의 주요 프로세스는 다음과 같습니다:

  1. 모든 바운딩 박스를 신뢰도에 따라 내림차순으로 정렬합니다.
  2. 가장 높은 신뢰도를 가진 박스를 선택하고, 나머지 박스와의 IoU(Intersection over Union)를 계산합니다.
  3. 사전에 정의된 임계치(threshold)를 기준으로 IoU가 threshold 이상의 박스를 제거합니다.
  4. 작업이 끝날 때까지 반복합니다.

NMS Python 예제 구현

아래는 Python을 사용하여 NMS를 구현하는 방법의 예시입니다. 먼저, 필요한 라이브러리를 설치해야 합니다:

pip install numpy
pip install opencv-python
        

이제 NMS 함수와 간단한 YOLO 예제를 작성해 보겠습니다:

import cv2
import numpy as np

def nms(boxes, scores, threshold):
    if len(boxes) == 0:
        return []

    # 박스의 좌표 분리
    boxes = np.array(boxes)
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]

    # 각 박스의 면적 계산
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)

    # 인덱스를 신뢰도에 따라 내림차순으로 정렬
    order = scores.argsort()[::-1]

    kept_indices = []
    
    while order.size > 0:
        i = order[0]  # 현재 가장 높은 점수를 가진 박스
        kept_indices.append(i)

        # IoU 계산
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        
        w = np.maximum(0, xx2 - xx1 + 1)
        h = np.maximum(0, yy2 - yy1 + 1)
        
        inter = w * h
        iou = inter / (areas[i] + areas[order[1:]] - inter)

        # Threshold를 기준으로 인덱스 제거
        order = order[np.where(iou <= threshold)[0] + 1]

    return kept_indices


# 테스트 데이터
boxes = [
    [50, 50, 150, 150],
    [60, 60, 160, 160],
    [70, 70, 170, 170],
]
scores = [0.9, 0.8, 0.7]

# NMS 실행
threshold = 0.3
kept_indices = nms(boxes, scores, threshold)
print("최종 선택된 박스 인덱스:", kept_indices)
        

위 코드는 주어진 바운딩 박스와 각 박스의 신뢰도를 입력으로 받아, NMS를 적용하여 최종적인 박스 인덱스를 반환합니다.

NMS의 유용성

NMS는 YOLO와 같은 객체 검출 시스템에서 필수적인 요소입니다. 객체가 겹치는 경우, NMS는 최종 결과에서 중복을 제거하여 더 정확한 예측을 제공합니다. 이는 모델의 성능을 향상시키고, 최종 사용자가 더 신뢰할 수 있는 예측 결과를 받을 수 있도록 합니다.

마무리

YOLO 모델을 사용한 객체 검출은 현대 컴퓨터 비전의 중요한 분야로 많은 사람들이 관심을 가지고 연구하고 있습니다. 본 글에서는 YOLO의 기본 원리와 비최대 억제(NMS)의 개념 및 실제 구현 방법에 대해 다루어 보았습니다. YOLO와 NMS를 활용한 다양한 프로젝트에서 여러분의 창의력을 발휘하시길 바랍니다.