def rotate_nms(rbboxes,
scores,
pre_max_size=None,
post_max_size=None,
iou_threshold=0.5,
flag=''):
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)
rbboxes = rbboxes[indices]
dets = torch.cat([rbboxes, 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(rotate_nms_cc(dets_np, iou_threshold), dtype=np.int64)
keep = ret[:post_max_size]
if keep.shape[0] == 0:
return torch.zeros([0]).long().to(rbboxes.device)
if pre_max_size is not None:
keep = torch.from_numpy(keep).long().to(rbboxes.device)
return indices[keep]
else:
return torch.from_numpy(keep).long().to(rbboxes.device)
def boxlist_rotate_nms(boxlist,
nms_thresh,
max_proposals=-1,
score_field="scores",
merge_nms=False):
"""
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)
"""
#nms_func = box_torch_ops.rotate_nms
rbboxes = boxlist.bbox[:, [0, 1, 3, 4, 6]]
#TODO try to multiply cls confidence
bev_scores = boxlist.get_field(score_field)[:, 0]
z_scores = boxlist.get_field(score_field)[:, 1]
dets = torch.cat([rbboxes, bev_scores.unsqueeze(-1)], dim=1)
dets_np = dets.data.cpu().numpy()
if len(dets_np) == 0:
keep = np.array([], dtype=np.int64)
boxlist = boxlist[keep]
return boxlist
if merge_nms == True:
ret = rotate_nms_merge(dets_np, nms_thresh)
merge_boxes_list = []
merge_score_list = []
for i in range(len(ret)):
# only merge z, h , cause box_score already choose best bev
zh_max_ids = boxlist[ret[i]].get_field("scores")[:, 1].max(0)[1]
merge_box_zh = boxlist[ret[i]].bbox[zh_max_ids, [2, 5]]
box_max_ids = boxlist[ret[i]].get_field("scores")[:, 0].max(0)[1]
merge_box_bev = boxlist[ret[i]].bbox[box_max_ids, [0, 1, 3, 4, 6]]
merge_box_stack = torch.cat(
(merge_box_bev[0][None], merge_box_bev[1][None],
merge_box_zh[0][None], merge_box_bev[2][None],
merge_box_bev[3][None], merge_box_zh[1][None],
merge_box_bev[4][None]))
merge_boxes_list.append(merge_box_stack)
merge_score_list.append(boxlist[ret[i]].get_field("scores").max())
merge_bboxes = torch.stack(merge_boxes_list)
merge_scores = torch.stack(merge_score_list)
boxlist = BoxList(merge_bboxes, (400, 352))
boxlist.add_field("scores", merge_scores)
return boxlist
ret = np.array(rotate_nms_cc(dets_np, nms_thresh), dtype=np.float32)
keep = ret[:max_proposals]
'''
merge rotate nms
return unmerge bboxlist and use boxlist.merge to merge the box
'''
#return empty boxlist
#if keep.shape[0] == 0:
# return torch.zeros([0]).long().to(rbboxes.device)
#TODO original return ids , now we return new bboxlist , so we have to write new Boxlist methed to create bboxlist from numpy
boxlist = boxlist[keep]
return boxlist