def find_top_rpn_proposals(self, rpn_cls_score_list, rpn_bbox_offset_list,
anchors_list, im_info):
prev_nms_top_n = (self.cfg.train_prev_nms_top_n
if self.training else self.cfg.test_prev_nms_top_n)
post_nms_top_n = (self.cfg.train_post_nms_top_n
if self.training else self.cfg.test_post_nms_top_n)
return_rois = []
for bid in range(im_info.shape[0]):
batch_proposal_list = []
batch_score_list = []
batch_level_list = []
for l, (rpn_cls_score, rpn_bbox_offset, anchors) in enumerate(
zip(rpn_cls_score_list, rpn_bbox_offset_list,
anchors_list)):
# get proposals and scores
offsets = rpn_bbox_offset[bid].transpose(2, 3, 0,
1).reshape(-1, 4)
proposals = self.box_coder.decode(anchors, offsets)
scores = rpn_cls_score[bid].transpose(1, 2, 0).flatten()
scores.detach()
# prev nms top n
scores, order = F.topk(scores,
descending=True,
k=prev_nms_top_n)
proposals = proposals[order, :]
batch_proposal_list.append(proposals)
batch_score_list.append(scores)
batch_level_list.append(F.full_like(scores, l))
# gather proposals, scores, level
proposals = F.concat(batch_proposal_list, axis=0)
scores = F.concat(batch_score_list, axis=0)
levels = F.concat(batch_level_list, axis=0)
proposals = layers.get_clipped_boxes(proposals, im_info[bid])
# filter invalid proposals and apply total level nms
keep_mask = layers.filter_boxes(proposals)
_, keep_inds = F.cond_take(keep_mask == 1, keep_mask)
proposals = proposals[keep_inds, :]
scores = scores[keep_inds]
levels = levels[keep_inds]
nms_keep_inds = layers.batched_nms(proposals, scores, levels,
self.cfg.rpn_nms_threshold,
post_nms_top_n)
# generate rois to rcnn head, rois shape (N, 5), info [batch_id, x1, y1, x2, y2]
rois = F.concat([proposals, scores.reshape(-1, 1)], axis=1)
rois = rois[nms_keep_inds]
batch_inds = F.full((rois.shape[0], 1), bid)
batch_rois = F.concat([batch_inds, rois[:, :4]], axis=1)
return_rois.append(batch_rois)
return_rois = F.concat(return_rois, axis=0)
return return_rois.detach()
def find_top_rpn_proposals(
self, rpn_bbox_offsets_list, rpn_cls_prob_list,
all_anchors_list, im_info
):
prev_nms_top_n = self.cfg.train_prev_nms_top_n \
if self.training else self.cfg.test_prev_nms_top_n
post_nms_top_n = self.cfg.train_post_nms_top_n \
if self.training else self.cfg.test_post_nms_top_n
batch_per_gpu = self.cfg.batch_per_gpu if self.training else 1
nms_threshold = self.cfg.rpn_nms_threshold
list_size = len(rpn_bbox_offsets_list)
return_rois = []
for bid in range(batch_per_gpu):
batch_proposals_list = []
batch_probs_list = []
batch_level_list = []
for l in range(list_size):
# get proposals and probs
offsets = rpn_bbox_offsets_list[l][bid].dimshuffle(2, 3, 0, 1).reshape(-1, 4)
all_anchors = all_anchors_list[l]
proposals = self.box_coder.decode(all_anchors, offsets)
probs = rpn_cls_prob_list[l][bid, 1].dimshuffle(1, 2, 0).reshape(1, -1)
# prev nms top n
probs, order = F.argsort(probs, descending=True)
num_proposals = F.minimum(probs.shapeof(1), prev_nms_top_n)
probs = probs.reshape(-1)[:num_proposals]
order = order.reshape(-1)[:num_proposals]
proposals = proposals.ai[order, :]
batch_proposals_list.append(proposals)
batch_probs_list.append(probs)
batch_level_list.append(mge.ones(probs.shapeof(0)) * l)
proposals = F.concat(batch_proposals_list, axis=0)
scores = F.concat(batch_probs_list, axis=0)
level = F.concat(batch_level_list, axis=0)
proposals = layers.get_clipped_box(proposals, im_info[bid, :])
# filter empty
keep_mask = layers.filter_boxes(proposals)
_, keep_inds = F.cond_take(keep_mask == 1, keep_mask)
proposals = proposals.ai[keep_inds, :]
scores = scores.ai[keep_inds]
level = level.ai[keep_inds]
# gather the proposals and probs
# sort nms by scores
scores, order = F.argsort(scores.reshape(1, -1), descending=True)
order = order.reshape(-1)
proposals = proposals.ai[order, :]
level = level.ai[order]
# apply total level nms
rois = F.concat([proposals, scores.reshape(-1, 1)], axis=1)
keep_inds = batched_nms(proposals, scores, level, nms_threshold, post_nms_top_n)
rois = rois.ai[keep_inds]
# rois shape (N, 5), info [batch_id, x1, y1, x2, y2]
batch_inds = mge.ones((rois.shapeof(0), 1)) * bid
batch_rois = F.concat([batch_inds, rois[:, :4]], axis=1)
return_rois.append(batch_rois)
return F.zero_grad(F.concat(return_rois, axis=0))