Merge NMS update
This commit is contained in:
parent
01ee0c5e95
commit
4a63b24b09
|
@ -555,7 +555,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, multi_label=T
|
|||
i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
|
||||
iou = box_iou(boxes, boxes[i]).tril_() # lower triangular iou matrix
|
||||
weights = (iou > iou_thres) * scores.view(-1, 1)
|
||||
weights /= weights.sum(0)
|
||||
weights /= weights.sum(0) + 1E-6
|
||||
x[i, :4] = torch.mm(weights.T, x[:, :4]) # merged_boxes(n,4) = weights(n,n) * boxes(n,4)
|
||||
elif method == 'vision':
|
||||
i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
|
||||
|
|
Loading…
Reference in New Issue