NMS高效实现
文章目录
- 01 什么是非极大值抑制
- 02 非极大值抑制基本实现
01 什么是非极大值抑制
Github: 非极大值抑制实现
非极大值抑制，简称为NMS算法，虽然在不同应用中实现的具体方式不太一样，但是思想还是一样的。下面我们看一张非极大值抑制的效果图：
左图是人脸检测的候选框结果，每个边界框都有一个置信度得分(Confidence Score)，如果不用非极大值抑制，那么就会有多个候选框出现。而右图则是经过极大值抑制之后的结果，其符合我们对人脸检测的预期效果。效果显然易见，如果不进行非极大值抑制，这结果简直没法看好嘛！下面我们来实现这个算法，包括相应的实现
02 非极大值抑制基本实现
首先，我们需要以下这些参数：
- 目标边界框的列表
- 对应的置信度列表
- 阀值
为了便于性能拓展，我们用类将其封装为抽象数据类型，将相应的属性和方法都封装在一起。现在我们先回顾一下非极大值抑制的基本流程：
while 直到原来置信度框中没有东西:
1. 根据置信度得分对其进行排序
2. 选择置信度最高的边界框，将其加入到最终输出的列表，然后我们将其从原来列表中删除
3. 计算 置信度最高的边界框 与 其它候选框 的 IoU
4. 删除IoU大于阀值的边界框
但是我们会做一个优化！我们将一些计算进行了矩阵化，事实上让我从零写也不一定写得出，调了好久。其实大家知道极大值抑制最原始思路是怎么做就可以了，要实现的时候来复制粘贴一下吧！
还有一点，numpy的算法效率真的很高，我尝试用numba.jit做了加速，发现相对于numpy来说还是慢了4-5倍。一般来说，最后能优化的就是多进程，速度大概还能提升4-5倍，但是既然这个模块不是瓶颈，其实没有什么必要去优化！都已经用Python了，既然性能不是瓶颈，那么就这样吧。真正要部署实际生产：
- 大多数时候我们在算法开发完成之后会用C++来重写。因为C++也有矩阵运算的包，也有OpenCv，因此大多数情况处理图像的时候，我更偏向于用OpenCv。
- 改到TF-Lite也不是很困难的事情，因为Numpy的基本操作TF都有。
在这里除了可以学习到如何绘制Bounding box，我们还能学习到怎样在原图上绘制Box，这是我们可以学习到的。
import numpy as np
import cv2
def nms(bounding_boxes, confidences, threshold):
"""
Args:
bounding_boxes: np.array([(x1, y1, x2, y2), ...])
confidences: np.array(conf1, conf2, ...),数量需要与bounding box一致,并且一一对应
threshold: IOU阀值,若两个bounding box的交并比大于该值，则置信度较小的box将会被抑制
Returns:
bounding_boxes: 经过NMS后的bounding boxes
confidences: 经过NMS后的confidences
"""
len_bound = bounding_boxes.shape[0]
len_conf = confidences.shape[0]
if len_bound != len_conf:
raise ValueError("Bounding box 与 Confidence 的数量不一致")
if len_bound == 0:
return np.array([]), np.array([])
bounding_boxes, confidences = bounding_boxes.astype(np.float), np.array(confidences)
x1, y1, x2, y2 = bounding_boxes[:, 0], bounding_boxes[:, 1], bounding_boxes[:, 2], bounding_boxes[:, 3]
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
idxs = np.argsort(confidences)
pick = []
while len(idxs) > 0:
# 因为idxs是从小到大排列的，last_idx相当于idxs最后一个位置的索引
last_idx = len(idxs) - 1
# 取出最大值在数组上的索引
max_value_idx = idxs[last_idx]
# 将这个添加到相应索引上
pick.append(max_value_idx)
xx1 = np.maximum(x1[max_value_idx], x1[idxs[: last_idx]])
yy1 = np.maximum(y1[max_value_idx], y1[idxs[: last_idx]])
xx2 = np.minimum(x2[max_value_idx], x2[idxs[: last_idx]])
yy2 = np.minimum(y2[max_value_idx], y2[idxs[: last_idx]])
w, h = np.maximum(0, xx2 - xx1 + 1), np.maximum(0, yy2 - yy1 + 1)
iou = w * h / areas[idxs[: last_idx]]
# 删除最大的value,并且删除iou > threshold的bounding boxes
idxs = np.delete(idxs, np.concatenate(([last_idx], np.where(iou > threshold)[0])))
# bounding box 返回一定要int类型,否则Opencv无法绘制
return np.array(bounding_boxes[pick, :]).astype(int), confidences[pick]
def plot_image_boxes(filename, bounding_boxes, confidences, title, color=(0, 0, 255)):
image = cv2.imread(filename)
boxes_number = confidences.shape[0]
for i in range(boxes_number):
(x1, y1, x2, y2), conf = bounding_boxes[i, :], confidences[i]
(w, h), baseline = cv2.getTextSize(str(conf),
fontFace=cv2.FONT_HERSHEY_SIMPLEX,
fontScale=1,
thickness=2)
cv2.rectangle(image, (x1, y1 - (2 * baseline + 5)), (x1 + w, y1), color, -1)
cv2.rectangle(image, (x1, y1), (x2, y2), color, 1)
cv2.putText(image, text=str(conf), org=(x1, y1), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
fontScale=1, color=(0, 0, 0), thickness=2)
cv2.imshow(title, image)
cv2.waitKey(0) # 按空格
if __name__ == '__main__':
bounding_boxes = np.array([(187, 82, 337, 317), (150, 67, 305, 282), (246, 121, 368, 304)], dtype=int)
confidences = np.array([0.9, 0.75, 0.8])
nms_boxes, nms_confs = nms(bounding_boxes, confidences, threshold=0.2)
plot_image_boxes('nms.jpg', bounding_boxes, confidences, title='Original')
plot_image_boxes('nms.jpg', nms_boxes, nms_confs, title='NMS')
