1. 程式人生 > >非極大值抑制NMS的python實現

非極大值抑制NMS的python實現

目標檢測中常用到NMS,在faster R-CNN中,每一個bounding box都有一個打分,NMS實現邏輯是:

1,按打分最高到最低將BBox排序 ,例如:A B C D E F

2,A的分數最高,保留。從B-E與A分別求重疊率IoU,假設B、D與A的IoU大於閾值,那麼B和D可以認為是重複標記去除

3,餘下C E F,重複前面兩步。

#coding:utf-8  
import numpy as np  
  
def py_cpu_nms(dets, thresh):  
    """Pure Python NMS baseline."""  
    x1 = dets[:, 0]  
    y1 = dets[:, 1]  
    x2 = dets[:, 2]  
    y2 = dets[:, 3]  
    scores = dets[:, 4]  #bbox打分
  
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)  
#打分從大到小排列,取index  
    order = scores.argsort()[::-1]  
#keep為最後保留的邊框  
    keep = []  
    while order.size > 0:  
#order[0]是當前分數最大的視窗,肯定保留  
        i = order[0]  
        keep.append(i)  
#計算視窗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:]])  
  
        w = np.maximum(0.0, xx2 - xx1 + 1)  
        h = np.maximum(0.0, yy2 - yy1 + 1)  
        inter = w * h  
#交/並得到iou值  
        ovr = inter / (areas[i] + areas[order[1:]] - inter)  
#inds為所有與視窗i的iou值小於threshold值的視窗的index,其他視窗此次都被視窗i吸收  
        inds = np.where(ovr <= thresh)[0]  
#order裡面只保留與視窗i交疊面積小於threshold的那些視窗,由於ovr長度比order長度少1(不包含i),所以inds+1對應到保留的視窗
        order = order[inds + 1]  
  
    return keep