NMS(非极大值抑制)实现

• 先看张图：

NMS算法思路来源于：https://chenzomi12.github.io/2016/12/14/YOLO-nms/

 

• 算法流程：

1. 把置信度最高的一个boundingbox(bbox)作为目标，然后对比剩下bbox与目标bbox之间的交叉区域
2. 如果交叉区域大于设定的阈值，那么在剩下的bbox中去除该bbox（即使该bbox的置信度与目标bbox的置信度一样）—-这个操作就是抑制最大重叠区域
3. 把第二置信度高的bbox作为目标，重复1、2

• v1版本：

import numpy as np


def nms(dets, thresh):
  x1 = dets[:, 0]
  y1 = dets[:, 1]
  x2 = dets[:, 2]
  y2 = dets[:, 3]
  score = dets[:, 4]
  
  sort_id_list = score.argsort()[::-1].tolist()
  res = []
  while len(sort_id_list) >= 1:
    i = sort_id_list.pop(0)
    res.append([x1[i], y1[i], x2[i], y2[i], score[i]])
    
    #intersect area left top point(xx1, yy1): xx1 >= x1, yy1 >= y1
    #intersect area right down point(xx2, yy2): xx2 <= x2, yy2 <= y2
    
    xx1 = np.maximum(x1[i], x1[sort_id_list[:]])
    yy1 = np.maximum(y1[i], y1[sort_id_list[:]])
    xx2 = np.minimum(x2[i], x2[sort_id_list[:]])
    yy2 = np.minimum(y2[i], y2[sort_id_list[:]])

    inter_w = np.maximum(0, (xx2 - xx1))
    inter_h = np.maximum(0, (yy2 - yy1))
    intersect = inter_w * inter_h
    
    #iou = intersect area / union; union = box1 + box2 - intersect
    iou = intersect / ((x2[i] - x1[i]) * (y2[i] - y1[i]) +\
    (x2[sort_id_list[:]] - x1[sort_id_list[:]]) * (y2[sort_id_list[:]] - y1[sort_id_list[:]]) -\
    intersect)
    for i in reversed(range(len(sort_id_list))):
      if iou[i] > thresh:
        sort_id_list.pop(i)

  return res
      




if __name__ == "__main__":
  
  dets = np.array([
                  [204, 102, 358, 250, 0.5],
                  [257, 118, 380, 250, 0.7],
                  [280, 135, 400, 250, 0.6],
                  [255, 118, 360, 235, 0.7]
                  ])


  thresh = 0.3
  res = nms(dets, thresh)
  print(res)

• v2版本：

import numpy as np


def nms(dets, thresh):
  x1 = dets[:, 0]
  y1 = dets[:, 1]
  x2 = dets[:, 2]
  y2 = dets[:, 3]
  score = dets[:, 4]
  
  order = score.argsort()[::-1]
  area = (x2 - x1) * (y2 - y1)
  res = []
  while order.size >= 1:
    i = order[0]
    res.append([x1[i], y1[i], x2[i], y2[i], score[i]])
    
    #intersect area left top point(xx1, yy1): xx1 >= x1, yy1 >= y1
    #intersect area right down point(xx2, yy2): xx2 <= x2, yy2 <= 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:]])

    w = np.maximum(0.0, (xx2 - xx1))
    h = np.maximum(0.0, (yy2 - yy1))
    intersect = w * h
    
    #iou = intersect area / union; union = box1 + box2 - intersect
    iou = intersect / (area[i] + area[order[1:]] - intersect)
    
    #update order index;ind +1:because ind is obtain by index [1:]
    ind = np.where(iou <= thresh)[0]
    order = order[ind +1]
    
  return res
      




if __name__ == "__main__":
  
  dets = np.array([
                  [204, 102, 358, 250, 0.5],
                  [257, 118, 380, 250, 0.7],
                  [280, 135, 400, 250, 0.6],
                  [255, 118, 360, 235, 0.7]
                  ])


  thresh = 0.7
  res = nms(dets, thresh)
  print(res)

• v1 vs v2 启示：

1、能在函数外通过一次计算实现的，尽量不要放到函数内进行多次循环计算，如area。

2、使用python高级库有助于快速实现。

3、np.array 数组本身可以作为另一个np.array数组的索引下标进行数组访问，如order[ind + 1](ind也为np.array数组)