计算机视觉 | 面试题：04、NMS详细工作机制及代码实现

问题

看到一句话：NMS都不懂，还做什么Detection !　虎躯一震……懂是大概懂，但代码能写出来吗？？？

在目标检测网络中，产生 proposal 后使用分类分支给出每个框的每类置信度，使用回归分支修正框的位置，最终会使用 NMS 方法去除同个类别当中 IOU 重叠度较高且 scores 即置信度较低的那些检测框。

下图就是在目标检测中 NMS 的使用效果：emmm大概就是能让你更无遮挡地看到美女的脸吧hhhh

背景知识

NMS (Non-maximum suppression) 非极大值抑制，即抑制不是极大值的检测框，根据什么去抑制？在目标检测领域，当然是根据 IOU (Intersection over Union) 去抑制。下图是绿色检测框与红色检测框的 IOU 计算方法：

NMS 原理及示例

注意 NMS 是针对一个特定的类别进行操作的。例如假设一张图中有要检测的目标有“人脸”和“猫”，没做NMS之前检测到10个目标框，每个目标框变量表示为: [ x 1 , y 1 , x 2 , y 2 , s c o r e 1 , s c o r e 2 ] [x_1,y_1,x_2,y_2,score_1,score_2] [x1​,y1​,x2​,y2​,score1​,score2​] ，其中 ( x 1 , y 1 ) (x_1,y_1) (x1​,y1​) 表示该框左上角坐标， ( x 2 , y 2 ) (x_2,y_2) (x2​,y2​) 表示该框右下角坐标， s c o r e 1 score_1 score1​ 表示"人脸"类别的置信度， s c o r e 2 score_2 score2​ 表示"猫"类别的置信度。当 s c o r e 1 score_1 score1​ 比 s c o r e 2 score_2 score2​ 大时，将该框归为“人脸”类别，反之归为“猫”类别。最后我们假设10个目标框中有6个被归类为“人脸”类别。

接下来演示如何对“人脸”类别的目标框进行 NMS 。

首先对6个目标框按照 s c o r e 1 score_1 score1​ 即置信度降序排序：

(1) 取出最大置信度的那个目标框 A 保存下来
(2) 分别判断 B-F 这５个目标框与 A 的重叠度 IOU ，如果 IOU 大于我们预设的阈值（一般为 0.5），则将该目标框丢弃。假设此时丢弃的是 C和 F 两个目标框，这时候该序列中只剩下 B D E 这三个。
(3) 重复以上流程，直至排序序列为空。

代码实现

# bboxees维度为 [N, 4]，scores维度为 [N, 1]，均为np.array()
def single_nms(self, bboxes, scores, thresh = 0.5):
    # x1、y1、x2、y2以及scores赋值
    x1 = bboxes[:, 0]
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]
    
    # 计算每个检测框的面积
    areas = (x2 - x1 + 1) * (y2 - y1 + 1) 
    
    # 按照 scores 置信度降序排序, order 为排序的索引
    order = scores.argsort() # argsort为python中的排序函数，默认升序排序
    order = order[::-1] # 将升序结果翻转为降序
    
    # 保留的结果框索引
    keep = []
    
    # torch.numel() 返回张量元素个数
    while order.size > 0:
        if order.size == 1:
            i = order[0]
            keep.append(i)
            break
        else:
            i = order[0]  # 在pytorch中使用item()来取出元素的实值，即若只是 i = order[0]，此时的 i 还是一个 tensor，因此不能赋值给 keep
            keep.append(i)
            
        # 计算相交区域的左上坐标及右下坐标
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        
        # 计算相交的面积，不重叠时为0
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        
        # 计算 IOU = 重叠面积 / (面积1 + 面积2 - 重叠面积)
        iou = inter / (areas[i] + areas[order[1:]] - inter)
        
        # 保留 IOU 小于阈值的 bboxes
        inds = np.where(iou <= thresh)[0]
        if inds.size == 0:
            break
        order = order[inds + 1] # 因为我们上面求iou的时候得到的结果索引与order相比偏移了一位，因此这里要补回来
    return keep  # 这里返回的是bboxes中的索引，根据这个索引就可以从bboxes中得到最终的检测框结果

参考资料

NMS算法详解（附Pytorch实现代码）
非极大值抑制（Non-Maximum Suppression，NMS）

标签：面试题,置信度,NMS,04,y1,score,np,order
来源： https://blog.csdn.net/Mrrunsen/article/details/120190382