NMS(非极大值抑制)算法 -- 理论、代码

目的

       消除 anchor-based 算法中因为滑动窗口或者重叠采样等方式产生的众多冗余候选框，保留效果最佳的框。

定义

       非极大值抑制（Non-Maximum Suppress，NMS）算法，就是抑制非极大值的目标（去冗余），从而搜索出局部极大值的目标（找最优）。

原理

输入

       给定所有可能的预测边框 predictions = [ [x_max, x_min, y_max, y_min, score], [*], ..., [*]] 以及一个给定的IoU阈值 iou_threshold

       有多少个类别，就执行多少次NMS算法

输出

       经过NMS算法过滤后的一类物体的预测框 result = [x_max, x_min, y_max, y_min, score].

算法步骤

（1）将所有可能的预测框按类别划分为num_class+1个集合，其中1为背景类，背景类无需NMS处理；

（2）对于每一个集合（类别），按类别分数从高到低进行排序，得到num_class个降序列表 list；

（3）从一个 list 中取得第一个元素（分数最高），逐个计算该元素与列表中剩余元素的IoU，若IoU大于给定阈值则将该元素从列表中删除，同时将第一个元素保留；

（4）对处理过后的降序列表 list 重复执行步骤（3），直至 list 为空；这样返回的keep列表中就是图中该类别所有的物体的唯一框。

（5）对每一个类别都执行步骤（3）~（4），直接遍历完所有的类别；

代码（添加注释、测试）

import cv2
import numpy as np
import matplotlib.pyplot as plt


# iou = np.array([0.57184716, 0.76505679])
# order = np.array([2, 1, 4, 5, 6, 8])
# inds = np.where(iou <= 0.8)[0]  # 将重叠度大于给定阈值的边框剔除掉，仅保留剩下的边框，返回相应的下标 [1,2]
# print('inds:', inds)  # 得到分数小于阈值的框索引 [0 1]
# order = order[inds + 1]  # 从剩余的候选框中继续筛选[1 2]
# print('order:', order)  # [1 4 ]


class NMS:
    def __init__(self, center=False, scale=1.0):
        """
        :param center: the format of coordinate -> diagonal [x1, y1, x2, y2] or center