From 23b34f4db8598d3916cdaa4c659baf9acfe01c3c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 25 Mar 2020 23:29:33 -0700 Subject: [PATCH] merge_batch NMS method --- utils/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/utils.py b/utils/utils.py index 134cf6ef..b7f11f56 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -561,7 +561,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, multi_label=T elif method == 'merge_batch': # Merge NMS i = torchvision.ops.boxes.nms(boxes, scores, iou_thres) iou = box_iou(boxes, boxes[i]).tril_() # upper triangular iou matrix - weights = (iou > conf_thres) * scores.view(-1, 1) + weights = (iou > iou_thres) * scores.view(-1, 1) weights /= weights.sum(0) pred[i, :4] = torch.matmul(weights.T, pred[:, :4]) # merged_boxes(n,4) = weights(n,n) * boxes(n,4) elif method == 'fast_batch': # FastNMS from https://github.com/dbolya/yolact