def nms(bboxes,
scores,
pre_max_size=None,
post_max_size=None,
iou_threshold=0.5):
if pre_max_size is not None:
num_keeped_scores = scores.shape[0]
pre_max_size = min(num_keeped_scores, pre_max_size)
scores, indices = torch.topk(scores, k=pre_max_size)
bboxes = bboxes[indices]
dets = torch.cat([bboxes, scores.unsqueeze(-1)], dim=1)
dets_np = dets.data.cpu().numpy()
if len(dets_np) == 0:
keep = np.array([], dtype=np.int64)
else:
ret = np.array(nms_gpu(dets_np, iou_threshold), dtype=np.int64)
keep = ret[:post_max_size]
if keep.shape[0] == 0:
return None
if pre_max_size is not None:
keep = torch.from_numpy(keep).long().cuda()
return indices[keep]
else:
return torch.from_numpy(keep).long().cuda()
def nms(bboxes,
scores,
pre_max_size=None, #1000
post_max_size=None, #300
iou_threshold=0.5):
if pre_max_size is not None:
num_keeped_scores = scores.shape[0]
pre_max_size = min(num_keeped_scores, pre_max_size)
indices = np.argsort(scores)[::-1][:pre_max_size] #top pre_max_size anchors boxes
scores = scores[indices]
bboxes = bboxes[indices] ## TODO: double check
dets = np.concatenate([bboxes, np.expand_dims(scores, axis=-1)], axis=1)
# dets_np = dets.data.cpu().numpy()
if len(dets) == 0:
keep = np.array([], dtype=np.int64)
else:
ret = np.array(nms_gpu(dets, iou_threshold), dtype=np.int64)
keep = ret[:post_max_size]
if keep.shape[0] == 0:
return None
if pre_max_size is not None:
return indices[keep]
else:
return keep