NMS(Non maximal supression) 계산하기, 파이썬

요약

---

Non maximal supression(NMS)은 예측된 오브젝트의 바운딩박스(bounding box)가 중복으로 여러 개가 할당 될 수 있을 때, 중복되는 바운딩박스를 제거하는 기술입니다. 기계적으로 무엇인가 학습하여 처리하는 것이 아닌, 사람이 고아낸 로직대로 제거하는 방식이기 때문에, 휴리스틱한(또는 hand-craft)한 방법으로도 논문에서 주로 소개가 됩니다. 

 

Figure 1. 한 오브젝트(포메라니안)에 다수의 바운딩박스가 예측된 예시. NMS은 이러한 여러박스 중에 최적의 박스남기는 작업이다.

 

NMS을 구현하는 방식은 아래와 같습니다.

1. 각 바운딩 박스별로 면적(area)을 구합니다.
2. 특정 박스를 하나 정하고, 정해진 박스외에 나머지 박스와의 오버렙 비율을 구합니다. 
   1.  두 박스의 x_min, y_min중에 최대를 구하면, 아래의 박스(//)의 좌상단 좌표를 구합니다.
   2. 두 박스의 x_max, y_max중에 최소를 구하면, 아래의 박스(//)의 우하단 좌표를 구합니다
3. 교집합의 박스와, 다른 박스들과의 오버렙 비율이 특정값(cutoff)이상을 넘기면, 1에서 지정한 박스를 삭제합니다(=인덱스를 제거합니다)
4. 이를 계속 반복합니다.

 

 

가령 아래와 같이 COCO 포맷(좌상단X, 좌상단Y, 가로, 세로)의 바운딩박스가 예측되었다고 생각합시다. 그러면, 첫 번째 [317, 86]의 바운딩박스에 대해서, 첫번째 바운딩 박스를 제외하고, 모든 바운딩박스와의 오버렙을 계산하면 됩니다.

// Array
[317, 86, 38, 41],
[464, 90, 40, 42],  # duplicated
[464, 90, 40, 42],  # duplicated

 

구체적인 방법은 아래와 같습니다.

def process_non_maximal_supression(
    boxes: np.ndarray, overlap_cutoff: float = 0.8
) -> np.ndarray:
    """NMS(Non maximal supression)을 시행함.

    Note:
        NMS 프로세스는 아래와 같다.
        1) 각 bounding box별로 Area(면적)을 계산함
        2) 특정 박스 하나에 대해서, 나머지박스와의 Overlap비율을 계산함
         - 각 박스는 (x_min, y_min, w, h)
         - 두 박스의 x_min, y_min중에 최대를 구하면, 아래의 박스(//)의 좌상단 좌표를 구함
         - 두 박스의 x_max, y_max중에 최소를 구하면, 아래의 박스(//)의 우하단 좌표를 구함
            |---------|
            |         |
            |----|------------|
            |    |////|       |
            |----|----|       |
                 |------------|

        3) 오버렙 컷오프(overlap_cutoff)가 주어진 cutoff보다 큰 경우, 2) 에서 선택된
        박스의 인덱스를 삭제
        4) 1~3)을 반복하여 남은 박스만 선택

    Args:
        boxes (np.ndarray): 각 array은 COCO Style box.
        overlap_cutoff (float, optional): _description_. Defaults to 0.7.

    Returns:
        np.ndarray: _description_
    """
    if len(boxes) == 0:
        return list()

    width = boxes[:, 2]
    height = boxes[:, 3]
    areas = width * height

    box_indices = np.arange(len(boxes))
    for box_idx, box in enumerate(boxes):

        # 자기자신을 제외한 다른 박스들 인덱스를 구함
        other_box_indices = box_indices[box_indices != box_idx]

        # 오버렙 계산을 위해 좌표를 구함
        xx1 = np.maximum(box[0], boxes[other_box_indices, 0])
        yy1 = np.maximum(box[1], boxes[other_box_indices, 1])
        xx2 = np.minimum(
            box[0] + box[2], boxes[other_box_indices, 0] + boxes[other_box_indices, 2]
        )
        yy2 = np.minimum(
            box[1] + box[3], boxes[other_box_indices, 1] + boxes[other_box_indices, 3]
        )

        w = np.maximum(0, xx2 - xx1 + 1)
        h = np.maximum(0, yy2 - yy1 + 1)

        overlap = (w * h) / areas[other_box_indices]
        if np.any(overlap > overlap_cutoff):
            box_indices = box_indices[box_indices != box_idx]

    return boxes[box_indices].astype(int)