merge NMS full matrix
This commit is contained in:
parent
8d788e10c4
commit
300e9a7ad6
|
@ -555,6 +555,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, multi_label=T
|
||||||
boxes, scores = x[:, :4].clone() + c.view(-1, 1) * max_wh, x[:, 4] # boxes (offset by class), scores
|
boxes, scores = x[:, :4].clone() + c.view(-1, 1) * max_wh, x[:, 4] # boxes (offset by class), scores
|
||||||
if method == 'merge': # Merge NMS (boxes merged using weighted mean)
|
if method == 'merge': # Merge NMS (boxes merged using weighted mean)
|
||||||
i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
|
i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
|
||||||
|
if n < 1E4: # update boxes
|
||||||
# weights = (box_iou(boxes, boxes).tril_() > iou_thres) * scores.view(-1, 1) # box weights
|
# weights = (box_iou(boxes, boxes).tril_() > iou_thres) * scores.view(-1, 1) # box weights
|
||||||
# weights /= weights.sum(0) # normalize
|
# weights /= weights.sum(0) # normalize
|
||||||
# x[:, :4] = torch.mm(weights.T, x[:, :4])
|
# x[:, :4] = torch.mm(weights.T, x[:, :4])
|
||||||
|
|
Loading…
Reference in New Issue