NMS 非極大值抑制 Python
阿新 • • 發佈:2018-12-12
#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