Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

merge_batch NMS method

This commit is contained in:
Glenn Jocher 2020-03-25 23:29:33 -07:00
parent aa0c64b5ac
commit 23b34f4db8
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
utils/utils.py
View File

@ -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 elif method == 'merge_batch': # Merge NMS
i = torchvision.ops.boxes.nms(boxes, scores, iou_thres) i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
iou = box_iou(boxes, boxes[i]).tril_() # upper triangular iou matrix 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) weights /= weights.sum(0)
pred[i, :4] = torch.matmul(weights.T, pred[:, :4]) # merged_boxes(n,4) = weights(n,n) * boxes(n,4) 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 elif method == 'fast_batch': # FastNMS from https://github.com/dbolya/yolact