The principle of NMS and its code realization

1. 为什么要用NMS

YOLOv3在预测阶段, Each target will generate at least 3个proposals, But a target generally only shows oneproposal, 因此需要对proposals进行去重,
这里去重的方法是NMS. 而NMSThe filter is based on IoU.

2. NMS的步骤

1. 先对所有proposalsDo confidence(confidence)的排序, Sort by descending order of confidence(从大到小排序).

2. will be of maximum confidenceproposal ($\mathrm{proposa}{\mathrm{l}}_{\max }$)取出来, 与剩下的proposals($\mathrm{proposal}{\mathrm{s}}_{\mathrm{rest}}$)进行IoU的计算, The purpose here is to filter the latterproposals, 意思是说:

   • 如果$\mathrm{proposal}{\mathrm{s}}_{\mathrm{rest}}$中的某一个proposal与$\mathrm{proposa}{\mathrm{l}}_{\max }$calculated betweenIoUvalue is less than the set threshold($\mathrm{thresh}$), 那么就认为这个proposal是可用的(retainable)

   • 一旦有IoU > $\mathrm{thresh}$的proposal, We think soproposal和$\mathrm{proposa}{\mathrm{l}}_{\max }$The same target is predicted(object), So the box is redundant(Because its confidence is not$\mathrm{proposa}{\mathrm{l}}_{\max }$of high confidence), 需要去除.

     the most confidentproposalwill definitely keep it, We are picking the restproposals

3. $\mathrm{proposa}{\mathrm{l}}_{\max }$与所有$\mathrm{proposal}{\mathrm{s}}_{\mathrm{rest}}$的IoUAfter calculating and filtering, The confidence pointer points to the next one(The second highest confidenceproposal)

4. 重复2, 3 -> 递归

NMSThe core main thing is: IoU > 阈值 的proposals需要去除, < 阈值的proposals则保留, The purpose is not to affect the forecast otherobject的proposals受到影响.

3. NMS的代码实现

import torch


def iou(proposal, proposals, isMin=False):
    """计算proposals的IoU 在计算IoU时, Need to find the intersection and union of the two. Suppose the coordinates of the two boxes are respectively: (x_11, y_11, x_12, y_12)和(x_21, y2_1, x_22, y_22) The coordinates of the intersection box: (max(x_11, x_21), max(y_11, y_21), min(x_12, x_22), min(y_12, y_22)) Args: proposal (_type_): 置信度最高的proposal -> [4] proposals (_type_): 剩余的proposals -> [N, 4] isMin (bool, optional): IoU的计算模式, 有两种: 1. (True) 交集 / 最小面积 2. (False -> Default) 交集 / 并集 Return: IoU (float): 返回proposal与proposals的IoU """
    # Calculates the area of ​​the current box: proposal = [x, y, w, h]
    box_area = (proposal[2] - proposal[0]) * (proposal[3] - proposal[1])
    
    # 计算proposalsThe area of ​​all boxes in  proposals = [N, [x, y, w, h]]
    boxes_area = (proposals[:, 2] - proposals[:, 0]) * (proposals[:, 3] - proposals[:, 1])

    # 计算交集proposal和proposals的计算
    xx_1 = torch.maximum(proposal[0], proposals[:, 0])  # The upper left corner of the intersectionx坐标
    yy_1 = torch.maximum(proposal[1], proposals[:, 1])  # The upper left corner of the intersectiony坐标
    xx_2 = torch.minimum(proposal[2], proposals[:, 2])  # Bottom right corner of the intersectionx坐标
    yy_2 = torch.minimum(proposal[3], proposals[:, 3])  # Bottom right corner of the intersectiony坐标
    
    # 特殊情况: The two boxes are not next to each other -> 没有交集
    w, h = torch.maximum(torch.Tensor([0]), xx_2 - xx_1), torch.maximum(torch.Tensor([0]), yy_2 - yy_1)
    
    # Get the area of ​​the intersecting boxes
    intersection_area = w * h
    
    if isMin:  # If one box is inside another box
        return intersection_area / torch.min(box_area, boxes_area)
    
    else:  # 两个框相交 -> 交集 / 并集
        return intersection_area / (box_area + boxes_area - intersection_area)
    
    
def nms(proposals, thresh=0.3, isMin=False):
    """Non-maximum suppression is used to remove redundancyproposals Args: proposals (torch.tensor): 网络推理得到的proposals -> [conf, x, y, w, h] thresh (float, optional): NMSFilter threshold. Defaults to 0.3. isMin (bool, optional): IoU的计算方式, Defaults to intersection/并集. Defaults to False. """
    
    # 根据proposalsThe confidences are sorted in descending order
    sorted_proposals = proposals[proposals[:, 0].argsort(descending=True)]

    # 定义一个ls, Used to save what needs to be preservedproposals
    keep_boxes = []
    
    while len(sorted_proposals) > 0:
        # 取出置信度最高的proposal并存放到ls中
        _box = sorted_proposals[0]  
        keep_boxes.append(_box)
        
        if len(sorted_proposals) > 1:
            # Take out the restproposals
            _boxes = sorted_proposals[1:]
            
            # 置信度最高的proposal与其他proposals进行IoU的计算
            """ 需要注意的是, NMSreserved during screeningIoU小于thresh的. 为什么? 两个proposal的IoU越小, 说明两个proposalThe framed objects are more different, 别忘了, NMS是为了去重, So need to keep less thanIoU的proposals torch.where(条件): 返回符合条件的索引 """
            sorted_proposals = _boxes[torch.where(iou(_box, _boxes, isMin) < thresh)]
        
        # 当剩下最后一个时候, 就不进行IoU计算了(Calculate with yourselfIoU没有意义)
        else:
            break
    
    # 将lsConvert to high dimensionaltensor
    return torch.stack(keep_boxes)


if __name__ == "__main__":
    proposal = torch.tensor(data=[0, 0, 4, 4])
    proposals = torch.tensor(data=[[4, 4, 5, 5],   # 没有交集
                                   [1, 1, 5, 5]])  # 有交集
    
    print(iou(proposal, proposals))  # tensor([0.0000, 0.3913])

    
    boxes = torch.tensor(data=[
                               [0.5, 1, 1, 10, 10],
                               [0.9, 1, 1, 11, 11],  # Very similar to the one above
                               [0.4, 8, 8, 12, 12]  # Not similar to the above two
    ])
    
    print(nms(boxes, thresh=0.1))
    """仅保留了2个 tensor([[ 0.9000, 1.0000, 1.0000, 11.0000, 11.0000], [ 0.4000, 8.0000, 8.0000, 12.0000, 12.0000]]) """
    
    print(nms(boxes, thresh=0.3))
    
    """All are reserved tensor([[ 0.9000, 1.0000, 1.0000, 11.0000, 11.0000], [ 0.5000, 1.0000, 1.0000, 10.0000, 10.0000], [ 0.4000, 8.0000, 8.0000, 12.0000, 12.0000]]) """
    ```