NMS非极大值抑制(Python)

NMS算法步骤:

对于某一类的所有预测框dets,都有其相对应的置信度score,设最终要输出的结果列表为keep=[ ]

  • 保存当前集合中置信度最高的边框至keep中,并从预测框集合中去掉它
  • 依次计算该预测框与其他所有边框的IOU
  • 去除所有IOU大于阈值的边框
  • 如果预测框集合不为空,返回第一步重复上述三个步骤。
def NMS(dets,thresh):
    #dets为某个类的框,形式为[[x1,y1,x2,y2,socre],[x1,y1,x2,y2,socre]]
    x1 = dets[:,0]
    y1 = dets[:,1]
    x2 = dets[:,2]
    y2 = dets[:,3]
    scores = dets[:4]
    # 每一个框的面积
    areas = (x2-x1+1)*(y2-y1+1)
    # 按照score降序排列,其中orders是序号
    orders = scores.argsort()[::-1]
    keep = []# 最终结果保存在keep中

    while orders.size()>0:
        i = orders[0]
        keep.append(i) # 保留该类剩余box中得分最高的一个

        #得到相交区域的左上和右下坐标
        xx1 = np.maximum(x1[i],x1[orders[1:]])
        yy1 = np.maximum(y1[i],y1[orders[1:]])
        xx2 = np.maximum(x2[i],x2[orders[1:]])
        yy2 = np.maximum(y2[i],y2[orders[1:]])

        w = np.maximum(0.0,xx2-xx1+1)
        h = np.maximum(0.0,yy2-yy1+1)
        # 计算出相交区域的面积,不重叠时面积为0
        inter = w*h
        over = inter/(areas[i]+areas[orders[1:]]-inter)
        # 保留IOU小于阈值的box
        inds = np.where(over<=thresh)[0]
        orders = orders[inds+1]
    return keep

 

你可能感兴趣的:(pytorch,机器学习,深度学习)