Non-Maximum Suppression
lastmod 2023-01-05created 2017-03-146 edits
About
Non-max suppression (non-max suppresion) is a technique for selecting the correct bounding box among the bounding boxes predicted by an object detector.
Objects in an image come in various sizes and shapes.
To detect them perfectly, object detection algorithms generate multiple bounding boxes.
Only one of these bounding boxes should be selected, and the technique applied for this is non-max suppression.
In the picture below, you can see that after Non-max suppresion is applied, only one bounding box remains for each object.
Non-maximum_suppression_in_yolo.png
Steps of the Non-max suppression algorithm
- From the list of bounding boxes for one class, select the bounding box with the highest score and remove it from the list. Then add it to the final boxes.
- Compare the selected bounding box with every bounding box in the bounding boxes list by computing the IoU. If the IoU is higher than the threshold, remove that box from the bounding boxes list.
- From the bounding boxes remaining in the list, select the one with the highest score and remove it from the list. Then add it to the final boxes.
- Again compare the IoU of the selected bounding box with the boxes in the list. If it is higher than the threshold, remove them from the list.
- Repeat until nothing remains in the bounding boxes.
- Repeat the above process for each class.
The link below is a yolo presentation by a Russian researcher that explains non-max suppression well.
Implementation
Implementing Non-max suppression in PyTorch
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