1. 程式人生 > >NMS 非極大值抑制 Python

NMS 非極大值抑制 Python

#coding:utf-8  
import numpy as np  
#nms 非極大值抑制
#輸入為bbox資料,包含[ymin,xmin,ymax,xmax];每個bbox的評分;每個bbox的標籤,如果沒有標籤就註釋掉
def py_cpu_nms(dets,socres,label, thresh=0.4):  
    x1 = dets[:, 0]  
    y1 = dets[:, 1]  
    x2 = dets[:, 2]  
    y2 = dets[:, 3]  
    scores = socres  #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:]])  
        
#         print xx1.shape
        
        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)  
#         print ovr
        inds=[]
#         print label[0] 
        b=order[0]
#inds為所有與視窗i的iou值小於threshold值的視窗的index,其他視窗此次都被視窗i吸收
#只有在標籤一致的情況下才會判斷是否被視窗吸收,如果沒有標籤的話就使用下面註釋的程式碼
        for j in range(len(order)-1):
            a=order[j+1]
            if label[a]==label[b]:
                if ovr[j]<=thresh:
                    inds.append(j+1) 
            else:
                inds.append(j+1)
#order裡面只保留與視窗i交疊面積小於threshold的那些視窗,由於ovr長度比order長度少1(不包含i),所以inds+1對應到保留的視窗
        order = order[inds]  
    
#           inds = np.where(ovr <= thresh)[0]  
#           order = order[inds + 1]  


    return keep