深度学习五 —— 手撕NMS

手撕NMS

概念

NMS就是非极大抑制，也没啥可说的，直接上代码

手撕NMS

import numpy as np
def NMS(dets, thresh):
    x1 = dets[:, 0]  # xmin
    y1 = dets[:, 1]
    x2 = dets[:, 2]  # xmax
    y2 = dets[:, 3]
    scores = dets[:, 4]  # confidence
 
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)  # 每一个候选框的面积
    order = np.argsort(scores[::-1])  # 按照score的降序排列
 
    keep = []  # 用来保存最后留下来的bounding box
    while order.size > 0:
        i = order[0]  # 置信度最高的bounding box的index
        keep.append(i)  # 添加本次置信度最高的bounding box的index
 
        # 当前bbox和剩下bbox之间的交叉区域
        # 选择大于x1,y1和小于x2,y2的区域
        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:]])  # 交叉区域右下角的纵坐标
 
        # 当前bbox和其他剩下bbox之间交叉区域的面积
        w = np.maximum(0., xx2 - xx1 + 1)
        h = np.maximum(0., yy2 - yy1 + 1)
        inter = w * h
 
        # 交叉区域面积 / (bbox + 某区域面积 - 交叉区域面积)
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        # 保留交集小于一定阈值的 bounding box
        inds = np.where(ovr <= thresh)[0]
        # 将order序列更新，由于前面得到的矩形框索引要比矩形框在原order序列中的索引小1，所以要把这个1加回来
        order = order[inds + 1]
 
    return keep
 
 
if __name__ == "__main__":
    dets = np.array([
        [200, 100, 380, 280, 0.5],
        [250, 150, 400, 300, 0.7],
        [280, 180, 420, 320, 0.6],
        [300, 200, 450, 350, 0.8]
    ])
    thread = 0.3
    res = NMS(dets, thread)
    pass