三维点云处理-NMS代码实现

三维点云处理-NMS代码实现

import numpy as np

def nms(dets,thresh):
  '''
  Python实现NMS
  输入：dets:[x1,y1,x2,y2,scores],分别为检测框的左上方和右下方角点坐标,以及对应的置信度得分
       thresh:IoU阈值，大于该阈值时会被去除
  '''
  # 分别获取检测框的角点坐标和置信度得分
  x1 = dets[:,0]
  y1 = dets[:,1]
  x2 = dets[:,2]
  y2 = dets[:,3]
  socres = dets[:,4]

  # 计算所有检测框的面积
  areas = (x2 - x1 + 1) * (y2 - y1 + 1)
  # 按照置信度从高到低排序，获取排序后的index
  order = scores.argsort()[::-1]
  # keep:保存最后保留的检测框的index，注意该index与order中的顺序一致
  keep = []
  # 循环进行NMS，直到order中存储的检测框为空
  while order.size > 0:
    # 每次循环开始时，order中存放的第一个检测框都对应于该类别中检测出的新的目标，直接将其存入keep中
    i = order[0]
    keep.append(i)
    # IoU计算过程
    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:]])

    w = mp.maximum(0.0,xx2-xx1+1)
    h = np.maximum(0.0,yy2-yy1+1)

    inter = w * h
    ovr = inter / (areas[i] + areas[order[1:]] - inter)
    
    # 选择IoU小于阈值的检测框进行保留用于下一轮循环
    inds = np.where(ovr <= thresh)
    order = order[inds+1]

  return keep