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