定义
在模型预测阶段，我们先为图像生成多个锚框，并为这些锚框一一预测类别和偏移量。随后，我们根据锚框及其预测偏移量得到预测边界框。当锚框数量较多时，同一个目标上可能会输出较多相似的预测边界框。了使结果更加简洁，我们可以移除相似的预测边界框。常用的方法叫作非极大值抑制（non-maximum suppression，NMS）。
计算流程
非极大值抑制的流程如下：
- 根据置信度得分进行排序
- 选择置信度最高的比边界框添加到最终输出列表中，将其从边界框列表中删除
- 计算所有边界框的面积
- 计算置信度最高的边界框与其它候选框的IoU。
- 删除IoU大于阈值的边界框
- 重复上述过程，直至边界框列表为空。
代码实现
下面的实现代码摘自《动手学深度学习（TF2.0版）》
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
| from collections import namedtuple
Pred_BB_Info = namedtuple("Pred_BB_Info",
["index", "class_id", "confidence", "xyxy"])
def non_max_suppression(bb_info_list, nms_threshold=0.5):
"""
非极大抑制处理预测的边界框
Args:
bb_info_list: Pred_BB_Info的列表, 包含预测类别、置信度等信息
nms_threshold: 阈值
Returns:
output: Pred_BB_Info的列表, 只保留过滤后的边界框信息
"""
output = []
sorted_bb_info_list = sorted(bb_info_list,
key = lambda x: x.confidence,
reverse=True)
while len(sorted_bb_info_list) != 0:
best = sorted_bb_info_list.pop(0)
output.append(best)
if len(sorted_bb_info_list) == 0:
break
bb_xyxy = []
for bb in sorted_bb_info_list:
bb_xyxy.append(bb.xyxy)
iou = compute_jaccard(tf.convert_to_tensor(best.xyxy),
tf.squeeze(tf.convert_to_tensor(bb_xyxy), axis=1))[0]
n = len(sorted_bb_info_list)
sorted_bb_info_list = [
sorted_bb_info_list[i] for i in
range(n) if iou[i] <= nms_threshold]
return output
|