非极大值抑制:
步骤如下(假设进行非极大值抑制的输入为2000x20的矩阵,2000表示该图像上框的个数,20表示类别数):
① 对2000×20维矩阵中每列按从大到小进行排序(每列表示一类,共20类。同一类可能有多个目标,如上图有两个人);
② 从每列最大的得分建议框开始,分别与该列后面的得分建议框进行IoU计算,若IoU>阈值,则剔除得分较小的建议框,否则认为图像中同一类物体有多个目标;[两个同类的目标的建议框基本不会有重叠,因此去掉建议框重叠较大的实际上是实现了剔除同一个目标的重叠框]
③ 从每列次大的得分建议框开始,重复步骤②;
④ 重复步骤③直到遍历完该列所有建议框;
⑤ 遍历完2000×20维矩阵所有列,即所有物体种类都做一遍非极大值抑制
import numpy as np
import random
import cv2
def non_max_suppress(predicts_dict, threshold):
for object_name, bbox in predicts_dict.items():#对每一个类别分别进行NMS;一次读取一对键值(即某个类别的所有框)
bbox_array = np.array(bbox, dtype=np.float)
#下面分别获取框的左上角坐标(x1,y1),右下角坐标(x2,y2)及此框的置信度;这里需要注意的是图像左上角可以看做坐标点(0,0),右下角可以看做坐标点(1,1),也就是说从左往右x值增大,从上往下y值增大
x1 = bbox_array[:, 0]
y1 = bbox_array[:, 1]
x2 = bbox_array[:, 2]
y2 = bbox_array[:, 3]
scores = bbox_array[:, 4]
#argsort函数返回的是数组值从小到大的索引值,[::-1]表示取反。即这里返回的是数组值从大到小的索引值
order = scores.argsort()[::-1]
#当前类所有框的面积(python会自动使用广播机制,相当于MATLAB中的.*即两矩阵对应元素相乘);x1=3,x2=5,习惯上计算x方向长度就是x=3、4、5这三个像素,即5-3+1=3,而不是5-3=2,所以需要加1
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
keep = []
#按confidence从高到低遍历bbx,移除所有与该矩形框的IoU值大于threshold的矩形框
while order.size > 0:
i = order[0]
keep.append(i)#保留当前最大confidence对应的bbx索引
#获取所有与当前bbx的交集对应的左上角和右下角坐标,并计算IoU(注意这里是同时计算一个bbx与其他所有bbx的IoU)
xx1 = np.maximum(x1[i], x1[order[1:]])
#最大置信度的左上角坐标分别与剩余所有的框的左上角坐标进行比较,分别保存较大值;因此这里的xx1的维数应该是当前类的框的个数减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:]])
inter = np.maximum(0.0, xx2-xx1+1) * np.maximum(0.0, yy2-yy1+1)
iou = inter / (areas[i] + areas[order[1:]] - inter)
#注意这里都是采用广播机制,同时计算了置信度最高的框与其余框的IoU
inds = np.where(iou <= threshold)[0]
#保留iou小于等于阙值的框的索引值
order = order[inds + 1]
#将order中的第inds+1处的值重新赋值给order;即更新保留下来的索引,加1是因为因为没有计算与自身的IOU,所以索引相差1,需要加上
bbox = bbox_array[keep]
predicts_dict[object_name] = bbox.tolist()
#predicts_dict = predicts_dict
return predicts_dict