[Object Detection] 비-최대 억제(NMS, Non-maximum Suppression)를 이해하고 파이토치로 구현하기

 안녕하세요! 이번 포스팅에서는 비-최대 억제(NMS,Non-maximum Suppression)을 알아보도록 하겠습니다.

 

 비최대 억제를 이해하기 위해서는 IoU(intersection over unio)에 대한 개념을 알아야합니다.

 IoU에 대한 내용은 아래 링크에 있습니다.

 

[Object Detection] IoU(Intersection over Union)를 이해하고 파이토치로 구현하기

 

비-최대 억제(MNS, Non-maximum Suppression)은 무엇일까요?

 비-최대 억제(non-max suppresion)은 object detector가 예측한 bounding box 중에서 정확한 bounding box를 선택하도록 하는 기법입니다.

 

 이미지에서 객체는 다양한 크기와 형태로 존재합니다.

 그리고 이것을 완벽하게 검출하기 위해 object detection 알고리즘은 여러 개의 bounding boxes를 생성합니다.

 이 중 하나의 bounding box만을 선택해야 하는데 이때 적용하는 기법이 non-max suppression 입니다.

 

 아래 사진을 보면 Non-max suppresion이 적용된 후 각각의 객체에 대한 bounding box가 1개만 남은 것을 확인할 수 있습니다.

 

 

Non-max suppression 알고리즘 작동 단계

1. 하나의 클래스에 대한 bounding boxes 목록에서 가장 높은 점수를 갖고 있는 bounding box를 선택하고 목록에서 제거합니다. 그리고 final box에 추가합니다.

 

2. 선택된 bounding box를 bounding boxes 목록에 있는 모든 bounding box와 IoU를 계산하여 비교합니다. IoU가 threshold보다 높으면 bounding boxes 목록에서 제거합니다.

 

3. bounding boxes 목록에 남아있는 bounding box에서 가장 높은 점수를 갖고 있는 것을 선택하고 목록에서 제거합니다. 그리고 final box에 추가합니다.

 

4. 다시 선택된 bounding box를 목록에 있는 box들과 IoU를 비교합니다. threshold보다 높으면 목록에서 제거합니다.

 

5. bounding boxes에 아무것도 남아 있지 않을 때 까지 반복합니다.

 

6. 각각의 클래스에 대해 위 과정을 반복합니다. 

 

 

아래 링크는 러시아의 연구원의 yolo 발표자료인데 non-max suppression에 대해 잘 나와있습니다.

docs.google.com/presentation/d/1aeRvtKG21KHdD5lg6Hgyhx5rPq_ZOsGjG5rJ1HP7BbA/pub?

 

Non-max suppression 파이토치로 구현하기

 위 알고리즘을 파이토치로 구현한 코드입니다.

 

import torch

from IoU import intersection_over_union

def nms(bboxes, iou_threshold, threshold, box_format='corners'):

    # bboxes가 list인지 확인합니다.
    assert type(bboxes) == list

    # box 점수가 threshold보다 높은 것을 선별합니다.
    # box shape는 [class, score, x1, y1, x2, y2] 입니다.
    bboxes = [box for box in bboxes if box[1] > threshold]
    # 점수 오름차순으로 정렬합니다.
    bboxes = sorted(bboxes, key=lambda x: x[1], reverse=True)
    bboxes_after_nmn = []

    # bboxes가 모두 제거될때 까지 반복합니다.
    while bboxes:
        # 0번째 index가 가장 높은 점수를 갖고있는 box입니다. 이것을 선택하고 bboxes에서 제거합니다.
        chosen_box = bboxes.pop(0)

        # box가 선택된 box와의 iou가 임계치보다 낮거나
        # class가 다르다면 bboxes에 남기고, 그 이외는 다 없앱니다.
        bboxes = [box for box in bboxes if box[0] != chosen_box[0] \
               or intersection_over_union(torch.tensor(chosen_box[2:]),
                                          torch.tensor(box[2:]),
                                          box_format=box_format)
                    < iou_threshold]

        # 선택된 박스를 추가합니다.
        bboxes_after_nmn.append(chosen_box)

    return bboxes_after_nmn

 

 다음 포스팅에서는 object detection의 평가 지표인 mAP에 대해 공부해보도록 하겠습니다.

 감사합니다.

 

---

출처

[1] towardsdatascience.com/non-maximum-suppression-nms-93ce178e177c

[2] www.analyticsvidhya.com/blog/2020/08/selecting-the-right-bounding-box-using-non-max-suppression-with-implementation/

[3] docs.google.com/presentation/d/1aeRvtKG21KHdD5lg6Hgyhx5rPq_ZOsGjG5rJ1HP7BbA/pub?start=false&loop=false&delayms=3000&slide=id.g137784ab86_4_4740

[4] www.youtube.com/watch?v=YDkjWEN8jNA