[NMS系列] Weighted-NMS

论文: openaccess.thecvf.com/content_ICC…

背景

NMS方法在一组候选框中选择分数最大的框作为最终的目标框。然而，非极大的检测结果保存了特征的最大值，因此忽略非极大的检测检测结果是不合适的。

那么与直接剔除机制不同，Weighted NMS顾名思义是对坐标加权平均，加权平均的对象包括 M 自身以及IoU≥NMS阈值的相邻框。

公式

代码

while detections.size(0):
    # 得到第一个bbx与其余bbx的iou大于nms_thres的判别(0, 1)， 1为大于，0为小于
    large_overlap = bbox_iou(detections[0, :4].unsqueeze(0), detections[:, :4]) > nms_thres

    # 判断他们的类别是否相同，只有相同时才进行nms， 相同时为1， 不同时为0
    label_match = detections[0, -1] == detections[:, -1]
    # Indices of boxes with lower confidence scores, large IOUs and matching labels

    # 只有在两个bbx的iou大于thres，且类别相同时，invalid为True，其余为False
    invalid = large_overlap & label_match

    # weights为对应的权值, 其格式为：将True bbx中的confidence连成一个Tensor
    weights = detections[invalid, 4:5]

    # Merge overlapping bboxes by order of confidence
    # 这里得到最后的bbx它是跟他满足IOU大于threshold，并且相同label的一些bbx，根据confidence重新加权得到
    # 并不是原始bbx的保留。

    detections[0, :4] = (weights * detections[invalid, :4]).sum(0) / weights.sum()
    keep_boxes += [detections[0]]

    # 去掉这些invalid
    detections = detections[~invalid]
    if keep_boxes:
        output[image_i] = torch.stack(keep_boxes)

这里做加权的时候，本身矩形框也按照置信度一起进行加权。