本篇文章主要是對NMS模塊的python實現,因為最近在做有關目標檢測的東西,而NMS作為目標檢測深度模型之后對圖片方框的后處理,是非常重要的一個部分。參考原文(NMS的python實現)
具體的原理建議參考原文,我這里就是搬運一下代碼, 當然也有增加了一些自己對代碼的解釋。
import numpy as np boxes=np.array([[100,100,210,210,0.72], [250,250,420,420,0.8], [220,220,320,330,0.92], [100,100,210,210,0.72], [230,240,325,330,0.81], [220,230,315,340,0.9]]) def py_cpu_nms(dets, thresh): #首先數據賦值和計算對應矩形框的面積
#dets的數據格式是dets=np.array[[xmin,ymin,xmax,ymax,scores]....]
x1 = dets[:,0] y1 = dets[:,1] x2 = dets[:,2] y2 = dets[:,3] areas = (y2-y1+1) * (x2-x1+1) scores = dets[:,4] print('areas ',areas) print('scores ',scores) #這邊的keep用於存放,NMS后剩余的方框
keep = [] # 取出分數從大到小排列的索引。.argsort()是從小到大排列,[::-1]是列表頭和尾顛倒一下。
index = scores.argsort()[::-1] print(scores) print(index) # 上面這兩句比如分數[0.72 0.8 0.92 0.72 0.81 0.9 ]
# 對應的索引index[ 2 5 4 1 3 0 ]記住是取出索引,scores列表沒變。scores[2]>scores[5]>scores[4]>...>scores[0]
# index會剔除遍歷過的方框,和合並過的方框。
while index.size >0: print(index.size) # 取出第一個方框進行和其他方框比對,看有沒有可以合並的
i = index[0] # every time the first is the biggst, and add it directly
#因為我們這邊分數已經按從大到小排列了。
#所以如果有合並存在,也是保留分數最高的這個,也就是我們現在那個這個
#keep保留的是索引值,不是具體的分數。
keep.append(i) # 計算交集的左上角和右下角
# 這里要注意,比如x1[i]這個方框的左上角x和所有其他的方框的左上角x的最大值
x11 = np.maximum(x1[i], x1[index[1:]]) # 返回值是【max(x1[i], x1[1]), max(x1[i], x1[2]),....】
y11 = np.maximum(y1[i], y1[index[1:]]) # np.maximum是逐個位置比較。
x22 = np.minimum(x2[i], x2[index[1:]]) y22 = np.minimum(y2[i], y2[index[1:]]) # 這邊要注意,如果兩個方框相交,X22-X11和Y22-Y11是正的。
# 如果兩個方框不相交,X22-X11和Y22-Y11是負的,我們把不相交的W和H設為0.
w = np.maximum(0, x22-x11+1) h = np.maximum(0, y22-y11+1) #計算重疊面積就是上面說的交集面積。不相交因為W和H都是0,所以不相交面積為0
overlaps = w*h print('overlaps is',overlaps) # 這個就是IOU公式(交並比)。
# 得出來的ious是一個列表,里面擁有當前方框和其他所有方框的IOU結果。
ious = overlaps / (areas[i]+areas[index[1:]] - overlaps) print('ious is',ious) # 接下來是合並重疊度最大的方框,也就是合並ious中值大於thresh的方框
# 我們合並的操作就是把他們剔除,因為我們合並這些方框只保留下分數最高的。
# 我們經過排序當前我們操作的方框就是分數最高的,所以我們剔除其他和當前重疊度最高的方框
# 這里np.where(ious<=thresh)[0]是一個固定寫法。
idx = np.where(ious <= thresh)[0] # 說明index[idx+1]這些位置是IOU小於指定閾值的位置
print('=============', idx) # 把留下來框再進行NMS操作
# 這邊留下的框是去除當前操作的框,和當前操作的框重疊度大於thresh的框
# 每一次都會先去除當前操作框(i框沒有參與計算IOU),所以索引的列表就會向前移動移位,要還原就+1,向后移動一位
index = index[idx+1] # because index start from 1
return keep print(py_cpu_nms(boxes, 0.5))
上面已經在關鍵地方進行了標注。大家可以下載下來試一試。