NMS(Non Max Suppression)

NMS(Non Max Suppression)

이번 포스팅 에서는 IOU에 이어서 NMS(Non Max Suppression)에 대해 알아보려고 한다. NMS는 여러 Object Detection논문(YOLO, SSD 등)에서 사용한 방법으로 각각의 물체에 대한 bounding box를 제일 잘 나타낸 1개의 box만 두고 나머지를 제거하는 기법이다. 자세히 알아보자.

모델이 위 그림과 같이 box를 예측했다고 해보자. 어떤 박스는 자동차를 거의 완벽하게 내포하고 있고 몇몇 박스는 조금씩 아쉽게 위치한 것을 볼 수 있다. 여기서 우리가 원하는 것은 이 박스들 중에 자동차를 가장 잘 포함시킨 박스 하나만 남겨서 깔끔하게 만드는 것이다. 이러한 작업은 어떻게 이루어 질까?

다시 박스를 바꿔서 모델이 위와 같은 예측을 했다고 생각해보자. 자동차라는 물체를 맞추기 위한 bounding box가 3개가 검출이 되었다 각 박스의 위에 써있는 숫자는 그 안에 물체가 있을 확률을 의미한다. 여기서 먼저 해야할 일은 positive 확률이 가장 높은 박스를 골라내는 것이다.

0.9의 positive 확률을 가진 박스가 메인 박스가 되면 이 메인박스와 그 주위에 있는 박스들과 IOU연산을 하게된다. 이때 그 값이 0.5가 넘어가면 해당 박스는 메인박스와 같은 물체를 예측하려 한다는 가정을 통해 해당 박스는 제거를 하게 된다. 그 과정은 위 그림과 같다. 

그런데 왜 굳이 주변에 있는 박스를 다 지우지 않고 IOU가 0.5이상인값만 지우는 것일까? 그 이유는 같은 물체가 서로 가까이 붙어있을때 문제가 발생하기 때문이다.

예를 들어 자동차 A와 자동차 B가 서로 가까이 붙어 있다고 해보자. 그중 자동차 A의 positive가 가장 높은 박스를 메인박스로 정하고 NMS를 수행하려고 한다. 그런데 메인박스 바로 옆에있는 박스가 자동차 B를 완벽하게 포함시키고 있는 박스인데 그냥 IOU를 계산하지 않고 주변 박스를 제거하게 된다면 옆에 있는 자동차 B를 잘 포함하는 박스까지 지워버리게 되는 것이다.

이러한 문제 때문에 메인박스와 주변 박스들간의 IOU를 통해 주변 박스가 메인박스와 같은 물체를 가리키는 것인지(IOU>0.5) 주변의 다른 물체의 박스를 가리키는 것인지(IOU<0.5)를 판단하는 것이다.

그런데 만약 사진에서 물체가 한종류가 아닌 여러 종류가 있다면 좀 더 고려해야 할 부분이 있다. NMS를 수행할 때는 반드시 각각의 종류에 맞는 bounding box끼리 진행해야 한다. 만약 자동차의 positive가 가장 높은 메인박스만 가지고 NMS를 진행하면 말을 예측하고 있는 모든 box들은 여전히 전부 살아있게 될 것이다.(말을 가리키는 박스들은 자동차의 메인박스와의 IOU가 0이기 때문에)

그러므로 모든 박스들은 positive확률과 함께 predicted class 정보까지 함께 가지고 있어야 하고 그 정보를 통해 각 class별로 구분지어 NMS를 진행해야 한다.

NMS를 정리해보면 다음과 같다.

1. 먼저 Positive확률이 낮은 박스는 바로 제거를 해준다 ex) 0.2 이하

2. 남은 box들에 대해 Positive가 높은 순으로 정렬한다.

3. 정렬된 box들을 위에서부터 하나씩 빼서 메인박스로 정하고 다른 박스들과 비교해서 같은 class인데 IOU가 0.5 이상인 박스는 제거를 한다.

---

구현코드

import torch
from iou import intersection_over_union

def non_max_suppression(bboxes, iou_threshold, prob_threshold, box_format='corners'):
    # bboxes = [[1, 0.9, x1, y1, x2, y2],[],[]...]
    #        = [class_prediction, probability_score, x1, y1, x2, y2]
    
    assert type(predictions) == list
    
    bboxes = [box for box in bboxes if box[1] > prob_threshold]
    bboxes = sorted(bboxes, key=lambda x: x[1], reverse=True)
    bboxes_after_nms = []
    
    while bboxes:
        chosen_box = bboxes.pop(0)
        
        bboxes = [box for box in bboxes 
                  if box[0] != chosen_box[0] 
                  or intersection_over_union(torch.Tensor(chosen_box[2:]).view(1,-1),
                                      torch.Tensor(box[2:]).view(1,-1),
                                      box_format=box_format) < iou_threshold]
        
        bboxes_after_nms.append(chosen_box)
        
    return bboxes_after_nms

IOU계산이 필요하기 떄문에 저번 포스팅에서 만든 intersection_over_union 함수를 불러왔다.

non_max_suppression 함수의 파라미터는

bboxes : [[class_prediction, probability_score, x1, y1, x2, y2], [], []]   -> 2차원 배열

iou_threshold : ex) 0.5   -> 박스제거 기준점

prob_threshold : ex) 0.3   -> NMS이전에 확률값이 작은 박스 제거

box_format : ‘corners’   -> 박스 좌표값 형태

이 코드에서는 박스를 한개씩 비교하기 때문에 intersection_over_union 함수에 넣을 좌표값을 view를 이용해 2차원으로 바꿔 주어야 한다.

---

End

이번 포스팅에서는 NMS(Non Max Suppression)에 대해 알아보았다. NMS는 출력을 하기 전에 같은 물체에 대해 겹치는 박스들을 제거하기 위해 쓰이는 기법이다. NMS에서 가장 중요한 포인트는 같은 종류에서 같은 물체를 가리키는 박스에 대해서만 제거를 해야 한다는 점이다. 그것을 우리는 IOU를 이용하여 판단을 하였다.

reference : www.youtube.com/watch?v=YDkjWEN8jNA&t=42s