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

Merge NMS update

This commit is contained in:
Glenn Jocher 2020-03-26 19:50:29 -07:00
parent 01ee0c5e95
commit 4a63b24b09
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

@ -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) i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
iou = box_iou(boxes, boxes[i]).tril_() # lower triangular iou matrix iou = box_iou(boxes, boxes[i]).tril_() # lower triangular iou matrix
weights = (iou > iou_thres) * scores.view(-1, 1) 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) x[i, :4] = torch.mm(weights.T, x[:, :4]) # merged_boxes(n,4) = weights(n,n) * boxes(n,4)
elif method == 'vision': elif method == 'vision':
i = torchvision.ops.boxes.nms(boxes, scores, iou_thres) i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)