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

问题

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

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

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

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

背景知识

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

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

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​ 即置信度降序排序:

目标框 score_1
A 0.9
B 0.85
C 0.7
D 0.6
E 0.4
F 0.1

(1) 取出最大置信度的那个目标框 A 保存下来
(2) 分别判断 B-F 这5个目标框与 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)

上一篇:记录一次腾讯实习投递经历(一)


下一篇:spring mvc 学习笔记【1】---前言