NMS, Soft-NMS 정리 및 구현

NMS (Non-Maximum Suppression)

Object Detction 모델이 객체를 정확하게 검출하기 위해 다양한 크기와 비율을 고려하여 하나의 이미지 안에 있는 여러 객체의 검출 값들(Label, Bounding Box, Score)을 구하게 됩니다.

이때, 모델은 하나의 객체에 대해 다양한 크기와 비율을 가진 여러개의 검출 값을 모두 사용하는 것은 비용적, 시각적으로 좋지 않습니다.

따라서, 여러개의 예측 값들 중에서 Label이 맞으며, Score가 가장 높고, 객체를 잘 표시하는 Bounding Box를 골라내야 하는데 여기서 Non-Maximun Suppression 알고리즘을 사용합니다.

https://pjreddie.com/darknet/yolov1/

즉, NMS 알고리즘은 Score가 가장 낮은 BBox(Bounding Box)를 억제(Suppress)하고 최상의(가장 적합한) BBox를 유지하기위해 사용합니다.

여기서 '억제(Suppress)'하기 위해서는 Score와 BBox의 IoU을 고려해야 합니다.

IoU (Intersection Over Union)
: BBox 간의 겹침 정도(Overlap)를 확인 할 수 있는 지표로서 객체 검출에 대한 성능 지표
Score
: Object Detection 모델이 BBox에 객체가 있음을 얼마나 확신하는지의 정도

NMS 알고리즘 과정

1. Score가 가장 높은 BBox 선택
2. 선택한 Box의 Overlap(IoU)를 다른 BBox와 비교
3. Overlap이 임계값(Threshold, 논문에선 0.5)값을 초과하는 BBox 제거
4. 다음으로 높은 Score의 BBox 선택
5. 2번~4번 과정 반복

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

Pytorch 구현

# NMS 계산
def non_max_suppression(bboxes, iou_threshold, threshold, box_format="midpoint"):
    """
    Does Non Max Suppression given bboxes

    Parameters:
        bboxes (list): list of lists containing all bboxes with each bboxes
        specified as [class_pred, prob_score, x1, y1, x2, y2]
        iou_threshold (float): threshold where predicted bboxes is correct
        threshold (float): threshold to remove predicted bboxes (independent of IoU) 
        box_format (str): "midpoint" or "corners" used to specify bboxes

    Returns:
        list: bboxes after performing NMS given a specific IoU threshold
    """

    assert type(bboxes) == list
	
    # bbox = [class, score, x1, y1, x2, y2]
    # bbox의 score가 threshold보다 높은 것을 선별
    bboxes = [box for box in bboxes if box[1] > threshold]
    # score 기준으로 오름차순 정렬
    bboxes = sorted(bboxes, key=lambda x: x[1], reverse=True)
    bboxes_after_nms = []
	
    while bboxes:
        chosen_box = bboxes.pop(0)

        # 다른 클래스의 bbox와 특정 threshold 이하의 박스를 남김
        # 다시말해서, 같은 클래스의 특정 threshold 이상의 박스를 없애는 것
        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_nms.append(chosen_box)

    return bboxes_after_nms

---

Soft-NMS

NMS 알고리즘의 과정에서 임계값(Threshold)를 설정하는 것은 매우 까다로우며 mAP가 낮아지는 문제가 있습니다.

그러한 이유는 동일한 클래스를 지닌 여러 물체가 뭉쳐있는 경우에, 하나의 바운딩 박스만을 검출하고, 나머지 바운딩 박스는 억제하기 때문입니다.

이를 해결하기 위해 Soft-NMS는 BBox를 억제(또는 제거)하는 것이 아니라, IoU의 값에 비례하여 Score를 감소시키는 방법을 사용하였습니다.

Soft-NMS 알고리즘 과정

기존 NMS는 다음과 같은 함수로 동작하며 이 함수에서 \(M\)은 해당 클래스에서 Score가 가장 높은 BBox를 의미하며, \(b_i\)는 동일한 클래스 내의 BBox를 의미합니다. 즉, \(M\)과 \(b_i\)의 IoU가 \(N_t\)값 이상이면, 0으로 억제됩니다.

NMS

Soft-NMS는 다음과 같은 함수로 동작합니다. \(M\)과 \(b_i\)의 IoU가 일정 값 이상일 때, 0으로 억제하는 것이 아닌, Score를 감소합니다. 이 함수는 \(N_t\) 입계값을 기준으로 Score 가 급격하게 변해, 연속형 함수가 아니라는 단점이 존재합니다.

Soft-NMS

이 단점을 해결하기위해 개선된 다음 함수를 제시합니다. 가우시안 분포를 활용하여 연속형 함수로 변환하여 부드러운 Score값을 출력할 수 있습니다.

Gaussian Soft-NMS

Pytorch 코드 구현

Github

GitHub - DocF/Soft-NMS: Python and Pytorch two implements of Soft NMS algorithm

def soft_nms_pytorch(dets, box_scores, sigma=0.5, thresh=0.001, cuda=0):
    """
    Build a pytorch implement of Soft NMS algorithm.
    # Augments
        dets:        boxes coordinate tensor (format:[y1, x1, y2, x2])
        box_scores:  box score tensors
        sigma:       variance of Gaussian function
        thresh:      score thresh
        cuda:        CUDA flag
    # Return
        the index of the selected boxes
    """

    # Indexes concatenate boxes with the last column
    N = dets.shape[0]
    if cuda:
        indexes = torch.arange(0, N, dtype=torch.float).cuda().view(N, 1)
    else:
        indexes = torch.arange(0, N, dtype=torch.float).view(N, 1)
    dets = torch.cat((dets, indexes), dim=1)

    # The order of boxes coordinate is [y1,x1,y2,x2]
    y1 = dets[:, 0]
    x1 = dets[:, 1]
    y2 = dets[:, 2]
    x2 = dets[:, 3]
    scores = box_scores
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)

    for i in range(N):
        # intermediate parameters for later parameters exchange
        tscore = scores[i].clone()
        pos = i + 1

        if i != N - 1:
            maxscore, maxpos = torch.max(scores[pos:], dim=0)
            if tscore < maxscore:
                dets[i], dets[maxpos.item() + i + 1] = dets[maxpos.item() + i + 1].clone(), dets[i].clone()
                scores[i], scores[maxpos.item() + i + 1] = scores[maxpos.item() + i + 1].clone(), scores[i].clone()
                areas[i], areas[maxpos + i + 1] = areas[maxpos + i + 1].clone(), areas[i].clone()

        # IoU calculate
        yy1 = np.maximum(dets[i, 0].to("cpu").numpy(), dets[pos:, 0].to("cpu").numpy())
        xx1 = np.maximum(dets[i, 1].to("cpu").numpy(), dets[pos:, 1].to("cpu").numpy())
        yy2 = np.minimum(dets[i, 2].to("cpu").numpy(), dets[pos:, 2].to("cpu").numpy())
        xx2 = np.minimum(dets[i, 3].to("cpu").numpy(), dets[pos:, 3].to("cpu").numpy())
        
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = torch.tensor(w * h).cuda() if cuda else torch.tensor(w * h)
        ovr = torch.div(inter, (areas[i] + areas[pos:] - inter))

        # Gaussian decay
        weight = torch.exp(-(ovr * ovr) / sigma)
        scores[pos:] = weight * scores[pos:]

    # select the boxes and keep the corresponding indexes
    keep = dets[:, 4][scores > thresh].int()

    return keep

 

NMS vs Soft-NMS

Soft-NMS가 NMS보다 약 1~2% 정도의 성능 향상이 있습니다.