NMS代码:
import tensorflow as tf
import numpy as np
def nms(boxes, overlap_threshold, mode='Union'): #输入boxes是一个N*5的numpy array,overlap_threshold是阈值
# if there are no boxes, return an empty list
if len(boxes) == 0:
return []
# if the bounding boxes integers, convert them to floats
if boxes.dtype.kind == "i":
boxes = boxes.astype("float")
# initialize the list of picked indexes
pick = []
# grab the coordinates of the bounding boxes
x1, y1, x2, y2, score = [boxes[:, i] for i in range(5)]
area = (x2 - x1 + 1) * (y2 - y1 + 1) #求每个bbox的面积
idxs = np.argsort(score) #从小到大排列
while len(idxs) > 0:
last = len(idxs) - 1 #注意:这句千万要注意,不能写成last = idxs[-1]!!!
i = idxs[last] #每次都从末尾取值,放入pick中
pick.append(i)
#计算前面每个框和当前最大框交集的左上角坐标,右下角坐标,不管有无交集,都可以得到这四个值
xx1 = np.maximum(x1[i], x1[idxs[:last]]) #拿当前最大的和前面所有的都比较一遍
yy1 = np.maximum(y1[i], y1[idxs[:last]])
xx2 = np.minimum(x2[i], x2[idxs[:last]])
yy2 = np.minimum(y2[i], y2[idxs[:last]])
# 计算bbox的宽度和高度,如果宽度或高度是负值,用0值代替
w = np.maximum(0, xx2 - xx1 + 1) #得到的w是一个长度为N的array,其中有一些值为0
h = np.maximum(0, yy2 - yy1 + 1) #h同理
inter = w * h #求每个iou的面积,得到的是一个长度为N的array,没有交集的框,inter为0
if mode == 'Min': #另一种求重叠的方法,这里没有用到
overlap = inter / np.minimum(area[i], area[idxs[:last]])
else:
overlap = inter / (area[i] + area[idxs[:last]] - inter) #交比并
# 将idxs中overlap满足阈值的bbox的index删除
# np.where(overlap > overlap_threshold)[0])得到的是长度为1的tuple的第一个值,是一个array,里面包含的是
# 满足这个条件的bbox的index,然后使用np.concatenate数组拼接函数,将原来score最大的那个bbox的index和现在
# 满足条件的bbox的index合并成一个numpy array,
# 最后使用delete删除掉指定bbox的index,这样idxs中就只剩下那些和之前最大score的bbox的overlap比较小的bbox
# 的index(重叠度大的那些已经删除了)
idxs = np.delete(idxs, np.concatenate(([last], np.where(overlap > overlap_threshold)[0])))
return pick
#-----------------------------------------
boxes = np.array([[10,10,40,80,0.5],[25,30,70,80,0.9],[30,25,65,78,0.8]])
overlap_threshold = 0.5
picks = nms(boxes,overlap_threshold,mode='Union')
print(picks)
参考这里 https://blog.csdn.net/wfei101/article/details/79918237?utm_source=blogxgwz1 解释的特别清晰,赞一个!
占个坑,下次来写iou代码~