图片来源
原理如下:
非极大值抑制的方法是:在进行目标检测时会产生很多候选区域,如图所示产生了1~5号框,而这些候选区域一般会有好多重叠的部分,先对这些框的概率从大到小进行排序,然后按概率从大到小遍历所有框:
(1)从最大概率1号矩形框开始,分别判断剩下的框与该框的重叠度IOU(IOU的计算)是否大于某个设定的阈值;
(2)假设2、3与1的IOU超过阈值,那么就扔掉2、3;并标记第一个矩形框1保留下来。
(3)从剩下的矩形框4、5中,选择概率最大的4,然后判断4、5的IOU,IOU大于一定的阈值,那么就扔掉;并标记4为第二个矩形框。
就这样一直重复,找到所有被保留下来的矩形框。
代码如下:
def non_max_suppression_fast(boxes, probs, overlap_thresh=0.9, max_boxes=300):
# code used from here: http://www.pyimagesearch.com/2015/02/16/faster-non-maximum-suppression-python/
# if there are no boxes, return an empty list
if len(boxes) == 0:
return []
# 获取框的坐标
x1 = boxes[:, 0]
y1 = boxes[:, 1]
x2 = boxes[:, 2]
y2 = boxes[:, 3]
np.testing.assert_array_less(x1, x2)
np.testing.assert_array_less(y1, y2)
# if the bounding boxes integers, convert them to floats --
# this is important since we'll be doing a bunch of divisions
if boxes.dtype.kind == "i":
boxes = boxes.astype("float")
# pick用来存放选取的框的索引
pick = []
# 计算框的面积
area = (x2 - x1) * (y2 - y1)
# 排序,返回框的概率从小到大的索引值
idxs = np.argsort(probs)
# 按框的概率从大到小循环遍历
while len(idxs) > 0:
# 获取索引列表中的最后一个索引,并将索引值添加到pick列表中
last = len(idxs) - 1
i = idxs[last]
pick.append(i)
# 计算选取的框与剩下框的交集
xx1_int = np.maximum(x1[i], x1[idxs[:last]])
yy1_int = np.maximum(y1[i], y1[idxs[:last]])
xx2_int = np.minimum(x2[i], x2[idxs[:last]])
yy2_int = np.minimum(y2[i], y2[idxs[:last]])
ww_int = np.maximum(0, xx2_int - xx1_int)
hh_int = np.maximum(0, yy2_int - yy1_int)
area_int = ww_int * hh_int
# 计算选取的框与剩下框的并集
area_union = area[i] + area[idxs[:last]] - area_int
# 计算交并比
overlap = area_int/(area_union + 1e-6)
# 删除掉重叠率较高的和最后一个索引(因为最后一个已经加入到pick中)
idxs = np.delete(idxs, np.concatenate(([last],
np.where(overlap > overlap_thresh)[0])))
#如果选取的框的数量达到上限就停止
if len(pick) >= max_boxes:
break
# return only the bounding boxes that were picked using the integer data type
boxes = boxes[pick].astype("int")
probs = probs[pick]
return boxes, probs