Spaces:
Runtime error
Runtime error
File size: 1,016 Bytes
8075387 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
from detectron2.layers import batched_nms
def ml_nms(boxlist, nms_thresh, max_proposals=-1,
score_field="scores", label_field="labels"):
"""
Performs non-maximum suppression on a boxlist, with scores specified
in a boxlist field via score_field.
Arguments:
boxlist(BoxList)
nms_thresh (float)
max_proposals (int): if > 0, then only the top max_proposals are kept
after non-maximum suppression
score_field (str)
"""
if nms_thresh <= 0:
return boxlist
if boxlist.has('pred_boxes'):
boxes = boxlist.pred_boxes.tensor
labels = boxlist.pred_classes
else:
boxes = boxlist.proposal_boxes.tensor
labels = boxlist.proposal_boxes.tensor.new_zeros(
len(boxlist.proposal_boxes.tensor))
scores = boxlist.scores
keep = batched_nms(boxes, scores, labels, nms_thresh)
if max_proposals > 0:
keep = keep[: max_proposals]
boxlist = boxlist[keep]
return boxlist
|