【程序】nms用来返回bboxes中的最大置信度框的下标：注意order[0] 和 i 是同一个变量

 

def nms(bboxes, scores, threshold=0.5):
    '''
    bboxes(tensor) [N,4]
    scores(tensor) [N,]
    '''
    x1 = bboxes[:, 0]
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]
    areas = (x2-x1) * (y2-y1)
    # print(scores)   # tensor([0.1006, 0.2381, 0.1185, 0.5342, 0.2892, 0.3521, 0.6027])
    _, order = scores.sort(0, descending=True)  # 降序排列score
    keep = []
    # print("order:", order)    # order: tensor([6, 3, 5, 4, 1, 2, 0])
    # print("order.numel:", order.numel())  # 7
    while order.numel() > 0:  # torch.numel()返回张量元素个数
        if order.numel() == 1:  # 保留框只剩一个
            # print("end1")
            # print(type(order))
            # print(order)
            i = order
            keep.append(i)
            break
        # print("len:", order.size())
        # print(keep)
        i = order[0]  # i=6，对应得分中最大的框   保留scores最大的那个框box[i]
        keep.append(i)

        # 计算box[i]与其余各框的IOU
        xx1 = x1[order[1:]].clamp(min=x1[i])  # [N-1,]
        yy1 = y1[order[1:]].clamp(min=y1[i])
        xx2 = x2[order[1:]].clamp(max=x2[i])
        yy2 = y2[order[1:]].clamp(max=y2[i])

        w = (xx2-xx1).clamp(min=0)
        h = (yy2-yy1).clamp(min=0)
        inter = w*h  # [N-1,]
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        ids = (ovr <= threshold).nonzero(as_tuple=False).squeeze()  # 注意此时idx为[N-1,] 而order为[N,]
        if ids.numel() == 0:
            # print("end2")
            break
        order = order[ids+1]   # 修补索引之间的差值
    # print(keep)
    return torch.LongTensor(keep)
    # return keep