def roi_pool(
rpn_fms,
rois,
stride,
pool_shape,
pooler_type="roi_align",
):
rois = rois.detach()
assert len(stride) == len(rpn_fms)
canonical_level = 4
canonical_box_size = 224
min_level = int(math.log2(stride[0]))
max_level = int(math.log2(stride[-1]))
num_fms = len(rpn_fms)
box_area = (rois[:, 3] - rois[:, 1]) * (rois[:, 4] - rois[:, 2])
assigned_level = F.floor(canonical_level +
F.log(F.sqrt(box_area) / canonical_box_size) /
np.log(2)).astype("int32")
assigned_level = F.minimum(assigned_level, max_level)
assigned_level = F.maximum(assigned_level, min_level)
assigned_level = assigned_level - min_level
# avoid empty assignment
assigned_level = F.concat([
assigned_level,
F.arange(num_fms, dtype="int32", device=assigned_level.device)
], )
rois = F.concat([rois, F.zeros((num_fms, rois.shape[-1]))])
pool_list, inds_list = [], []
for i in range(num_fms):
_, inds = F.cond_take(assigned_level == i, assigned_level)
level_rois = rois[inds]
if pooler_type == "roi_pool":
pool_fm = F.nn.roi_pooling(rpn_fms[i],
level_rois,
pool_shape,
mode="max",
scale=1.0 / stride[i])
elif pooler_type == "roi_align":
pool_fm = F.nn.roi_align(
rpn_fms[i],
level_rois,
pool_shape,
mode="average",
spatial_scale=1.0 / stride[i],
sample_points=2,
aligned=True,
)
pool_list.append(pool_fm)
inds_list.append(inds)
fm_order = F.argsort(F.concat(inds_list, axis=0))
pool_feature = F.concat(pool_list, axis=0)
pool_feature = pool_feature[fm_order][:-num_fms]
return pool_feature
def roi_pool(
rpn_fms,
rois,
stride,
pool_shape,
roi_type="roi_align",
):
assert len(stride) == len(rpn_fms)
canonical_level = 4
canonical_box_size = 224
min_level = math.log2(stride[0])
max_level = math.log2(stride[-1])
num_fms = len(rpn_fms)
box_area = (rois[:, 3] - rois[:, 1]) * (rois[:, 4] - rois[:, 2])
level_assignments = F.floor(canonical_level +
F.log(box_area.sqrt() / canonical_box_size) /
np.log(2))
level_assignments = F.minimum(level_assignments, max_level)
level_assignments = F.maximum(level_assignments, min_level)
level_assignments = level_assignments - min_level
# avoid empty assignment
level_assignments = F.concat(
[level_assignments,
mge.tensor(np.arange(num_fms, dtype=np.int32))], )
rois = F.concat([rois, mge.zeros((num_fms, rois.shapeof(-1)))])
pool_list, inds_list = [], []
for i in range(num_fms):
mask = level_assignments == i
_, inds = F.cond_take(mask == 1, mask)
level_rois = rois.ai[inds]
if roi_type == "roi_pool":
pool_fm = F.roi_pooling(rpn_fms[i],
level_rois,
pool_shape,
mode="max",
scale=1.0 / stride[i])
elif roi_type == "roi_align":
pool_fm = F.roi_align(
rpn_fms[i],
level_rois,
pool_shape,
mode="average",
spatial_scale=1.0 / stride[i],
sample_points=2,
aligned=True,
)
pool_list.append(pool_fm)
inds_list.append(inds)
fm_order = F.concat(inds_list, axis=0)
fm_order = F.argsort(fm_order.reshape(1, -1))[1].reshape(-1)
pool_feature = F.concat(pool_list, axis=0)
pool_feature = pool_feature.ai[fm_order][:-num_fms]
return pool_feature
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))
def fwd(data):
return F.argsort(data, True)
def find_top_rpn_proposals(is_train, rpn_bbox_offsets_list, rpn_cls_prob_list,
all_anchors_list, im_info):
prev_nms_top_n = config.train_prev_nms_top_n \
if is_train else config.test_prev_nms_top_n
post_nms_top_n = config.train_post_nms_top_n \
if is_train else config.test_post_nms_top_n
batch_per_gpu = config.batch_per_gpu if is_train else 1
nms_threshold = config.rpn_nms_threshold
box_min_size = config.rpn_min_box_size
bbox_normalize_targets = config.rpn_bbox_normalize_targets
bbox_normalize_means = config.bbox_normalize_means
bbox_normalize_stds = config.bbox_normalize_stds
list_size = len(rpn_bbox_offsets_list)
return_rois, return_probs = [], []
batch_per_gpu = rpn_cls_prob_list[0].shape[0]
for bid in range(batch_per_gpu):
batch_proposals_list = []
batch_probs_list = []
for l in range(list_size):
# get proposals and probs
offsets = rpn_bbox_offsets_list[l][bid] \
.transpose(1, 2, 0).reshape(-1, 4)
if bbox_normalize_targets:
std_opr = tensor(config.bbox_normalize_stds[None, :])
mean_opr = tensor(config.bbox_normalize_means[None, :])
pred_offsets = pred_offsets * std_opr
pred_offsets = pred_offsets + mean_opr
all_anchors = all_anchors_list[l]
proposals = bbox_transform_inv_opr(all_anchors, offsets)
if config.anchor_within_border:
proposals = clip_boxes_opr(proposals, im_info[bid, :])
probs = rpn_cls_prob_list[l][bid] \
.transpose(1,2,0).reshape(-1, 2)
probs = F.softmax(probs)[:, 1]
# gather the proposals and probs
batch_proposals_list.append(proposals)
batch_probs_list.append(probs)
batch_proposals = F.concat(batch_proposals_list, axis=0)
batch_probs = F.concat(batch_probs_list, axis=0)
# filter the boxes with small size.
wh = batch_proposals[:, 2:4] - batch_proposals[:, :2] + 1
thresh = box_min_size * im_info[bid, 2]
keep_mask = F.prod((wh >= thresh), axis=1)
keep_mask = keep_mask + F.equal(keep_mask.sum(), 0)
keep_mask, inds = F.cond_take(keep_mask > 0, keep_mask)
inds = inds.astype(np.int32)
# batch_proposals = F.nn.indexing_one_hot(batch_proposals, inds, 0)
# batch_probs = F.nn.indexing_one_hot(batch_probs, inds, 0)
batch_proposals, batch_probs = batch_proposals[inds], batch_probs[inds]
# prev_nms_top_n
num_proposals = F.minimum(prev_nms_top_n, batch_proposals.shape[0])
idx = F.argsort(batch_probs, descending=True)
topk_idx = idx[:num_proposals].reshape(-1)
batch_proposals = batch_proposals[topk_idx].detach()
batch_probs = batch_probs[topk_idx].detach()
# For each image, run a total-level NMS, and choose topk results.
keep_inds = nms(batch_proposals,
batch_probs,
nms_threshold,
max_output=2000)
# num = F.minimum(post_nms_top_n, keep_inds.shape[0])
# keep_inds = keep_inds[:num]
batch_rois, batch_probs = batch_proposals[keep_inds], batch_probs[
keep_inds]
# cons the rois
batch_inds = F.ones((batch_rois.shape[0], 1)) * bid
batch_rois = F.concat([batch_inds, batch_rois[:, :4]], axis=1)
return_rois.append(batch_rois)
return_probs.append(batch_probs)
if batch_per_gpu == 1:
return batch_rois, batch_probs
else:
concated_rois = F.concat(return_rois, axis=0)
concated_probs = F.concat(return_probs, axis=0)
return concated_rois, concated_probs
def cascade_roi_target(rpn_rois, im_info, gt_boxes, pos_threshold=0.5, top_k=1):
return_rois = []
return_labels = []
return_bbox_targets = []
# get per image proposals and gt_boxes
for bid in range(config.batch_per_gpu):
gt_boxes_perimg = gt_boxes[bid, :im_info[bid, 5], :]
batch_inds = mge.ones((gt_boxes_perimg.shapeof()[0], 1)) * bid
#if config.proposal_append_gt:
gt_rois = F.concat([batch_inds, gt_boxes_perimg[:, :4]], axis=1)
batch_roi_mask = rpn_rois[:, 0] == bid
batch_roi_inds = mask_to_inds(batch_roi_mask)
all_rois = F.concat([rpn_rois.ai[batch_roi_inds], gt_rois], axis=0)
overlaps_normal, overlaps_ignore = box_overlap_ignore_opr(
all_rois[:, 1:5], gt_boxes_perimg)
overlaps_normal, overlaps_normal_indices = F.argsort(overlaps_normal, descending=True)
overlaps_ignore, overlaps_ignore_indices = F.argsort(overlaps_ignore, descending=True)
# gt max and indices, ignore max and indices
max_overlaps_normal = overlaps_normal[:, :top_k].reshape(-1)
gt_assignment_normal = overlaps_normal_indices[:, :top_k].reshape(-1)
max_overlaps_ignore = overlaps_ignore[:, :top_k].reshape(-1)
gt_assignment_ignore = overlaps_ignore_indices[:, :top_k].reshape(-1)
# cons masks
ignore_assign_mask = (max_overlaps_normal < config.fg_threshold) * (
max_overlaps_ignore > max_overlaps_normal)
max_overlaps = max_overlaps_normal * (1 - ignore_assign_mask) + \
max_overlaps_ignore * ignore_assign_mask
gt_assignment = gt_assignment_normal * (1- ignore_assign_mask) + \
gt_assignment_ignore * ignore_assign_mask
gt_assignment = gt_assignment.astype(np.int32)
labels = gt_boxes_perimg.ai[gt_assignment, 4]
fg_mask = (max_overlaps >= config.fg_threshold) * (1 - F.equal(labels, config.ignore_label))
bg_mask = (max_overlaps < config.bg_threshold_high) * (
max_overlaps >= config.bg_threshold_low)
fg_mask = fg_mask.reshape(-1, top_k)
bg_mask = bg_mask.reshape(-1, top_k)
#pos_max = config.num_rois * config.fg_ratio
#fg_inds_mask = _bernoulli_sample_masks(fg_mask[:, 0], pos_max, 1)
#neg_max = config.num_rois - fg_inds_mask.sum()
#bg_inds_mask = _bernoulli_sample_masks(bg_mask[:, 0], neg_max, 1)
labels = labels * fg_mask.reshape(-1)
#keep_mask = fg_inds_mask + bg_inds_mask
#keep_inds = mask_to_inds(keep_mask)
#keep_inds = keep_inds[:F.minimum(config.num_rois, keep_inds.shapeof()[0])]
# labels
labels = labels.reshape(-1, top_k)
gt_assignment = gt_assignment.reshape(-1, top_k).reshape(-1)
target_boxes = gt_boxes_perimg.ai[gt_assignment, :4]
#rois = all_rois.ai[keep_inds]
target_shape = (all_rois.shapeof()[0], top_k, all_rois.shapeof()[-1])
target_rois = F.add_axis(all_rois, 1).broadcast(target_shape).reshape(-1, all_rois.shapeof()[-1])
bbox_targets = bbox_transform_opr(target_rois[:, 1:5], target_boxes)
if config.rcnn_bbox_normalize_targets:
std_opr = mge.tensor(config.bbox_normalize_stds[None, :])
mean_opr = mge.tensor(config.bbox_normalize_means[None, :])
minus_opr = mean_opr / std_opr
bbox_targets = bbox_targets / std_opr - minus_opr
bbox_targets = bbox_targets.reshape(-1, top_k * 4)
return_rois.append(all_rois)
return_labels.append(labels)
return_bbox_targets.append(bbox_targets)
if config.batch_per_gpu == 1:
return F.zero_grad(all_rois), F.zero_grad(labels), F.zero_grad(bbox_targets)
else:
return_rois = F.concat(return_rois, axis=0)
return_labels = F.concat(return_labels, axis=0)
return_bbox_targets = F.concat(return_bbox_targets, axis=0)
return F.zero_grad(return_rois), F.zero_grad(return_labels), F.zero_grad(return_bbox_targets)
def find_top_rpn_proposals(is_train, rpn_bbox_offsets_list, rpn_cls_prob_list,
all_anchors_list, im_info):
prev_nms_top_n = config.train_prev_nms_top_n \
if is_train else config.test_prev_nms_top_n
post_nms_top_n = config.train_post_nms_top_n \
if is_train else config.test_post_nms_top_n
batch_per_gpu = config.batch_per_gpu if is_train else 1
nms_threshold = config.rpn_nms_threshold
box_min_size = config.rpn_min_box_size
bbox_normalize_targets = config.rpn_bbox_normalize_targets
bbox_normalize_means = config.bbox_normalize_means
bbox_normalize_stds = config.bbox_normalize_stds
list_size = len(rpn_bbox_offsets_list)
return_rois = []
return_probs = []
for bid in range(batch_per_gpu):
batch_proposals_list = []
batch_probs_list = []
for l in range(list_size):
# get proposals and probs
offsets = rpn_bbox_offsets_list[l][bid] \
.dimshuffle(1, 2, 0).reshape(-1, 4)
if bbox_normalize_targets:
std_opr = tensor(config.bbox_normalize_stds[None, :])
mean_opr = tensor(config.bbox_normalize_means[None, :])
pred_offsets = pred_offsets * std_opr
pred_offsets = pred_offsets + mean_opr
all_anchors = all_anchors_list[l]
proposals = bbox_transform_inv_opr(all_anchors, offsets)
if config.anchor_within_border:
proposals = clip_boxes_opr(proposals, im_info[bid, :])
probs = rpn_cls_prob_list[l][bid] \
.dimshuffle(1,2,0).reshape(-1, 2)
probs = F.softmax(probs)[:, 1]
# gather the proposals and probs
batch_proposals_list.append(proposals)
batch_probs_list.append(probs)
batch_proposals = F.concat(batch_proposals_list, axis=0)
batch_probs = F.concat(batch_probs_list, axis=0)
# filter the zero boxes.
batch_keep_mask = filter_boxes_opr(batch_proposals,
box_min_size * im_info[bid, 2])
batch_probs = batch_probs * batch_keep_mask
# prev_nms_top_n
num_proposals = F.minimum(prev_nms_top_n, batch_probs.shapeof()[0])
batch_probs, idx = F.argsort(batch_probs, descending=True)
batch_probs = batch_probs[:num_proposals].reshape(-1, 1)
topk_idx = idx[:num_proposals].reshape(-1)
batch_proposals = batch_proposals.ai[topk_idx]
batch_rois = F.concat([batch_proposals, batch_probs], axis=1)
# For each image, run a total-level NMS, and choose topk results.
keep_inds = gpu_nms(batch_rois, nms_threshold, post_nms_top_n)
batch_rois = batch_rois.ai[keep_inds]
batch_probs = batch_rois[:, -1]
# cons the rois
batch_inds = mge.ones((batch_rois.shapeof()[0], 1)) * bid
batch_rois = F.concat([batch_inds, batch_rois[:, :-1]], axis=1)
return_rois.append(batch_rois)
return_probs.append(batch_probs)
if batch_per_gpu == 1:
return batch_rois, batch_probs
else:
concated_rois = F.concat(return_rois, axis=0)
concated_probs = F.concat(return_probs, axis=0)
return concated_rois, concated_probs
def fpn_roi_target(rpn_rois,
im_info,
gt_boxes,
fg_threshold=config.fg_threshold,
top_k=1):
return_rois, return_labels = [], []
return_bbox_targets = []
# get per image proposals and gt_boxes
batch_per_gpu = im_info.shape[0]
sampling = True
# is_sample = True if top_k < 2 else False
for bid in range(batch_per_gpu):
gt_boxes_perimg = gt_boxes[bid, :im_info[bid, 5].astype(np.int32), :]
dummy_gt = F.ones([1, gt_boxes_perimg.shape[1]])
batch_inds = F.ones((gt_boxes_perimg.shape[0], 1)) * bid
#if config.proposal_append_gt:
gt_rois = F.concat([batch_inds, gt_boxes_perimg[:, :4]], axis=1)
batch_rois_mask = F.equal(rpn_rois[:, 0], bid) > 0
_, batch_rois_index = F.cond_take(batch_rois_mask, batch_rois_mask)
# batch_roi_mask = rpn_rois[:, 0] == bid
# batch_roi_inds = mask_to_inds(batch_roi_mask)
all_rois= F.concat([rpn_rois[batch_rois_index], gt_rois], axis=0) if sampling \
else rpn_rois[batch_rois_index]
# all_rois = F.concat([rpn_rois.ai[batch_roi_inds], gt_rois], axis=0)
gt_boxes_perimg = F.concat([gt_boxes_perimg, dummy_gt], axis=0)
overlaps_normal, overlaps_ignore = box_overlap_ignore_opr(
all_rois[:, 1:5], gt_boxes_perimg)
# overlaps_normal, overlaps_normal_indices = F.argsort(overlaps_normal, descending=True)
# overlaps_ignore, overlaps_ignore_indices = F.argsort(overlaps_ignore, descending=True)
overlaps_normal_indices = F.argsort(overlaps_normal, descending=True)
overlaps_normal = F.gather(overlaps_normal, 1, overlaps_normal_indices)
# overlaps_normal = F.nn.indexing_one_hot(overlaps_normal, overlaps_normal_indices, 1)
overlaps_ignore_indices = F.argsort(overlaps_ignore, descending=True)
overlaps_ignore = F.gather(overlaps_ignore, 1, overlaps_ignore_indices)
# overlaps_ignore = F.nn.indexing_one_hot(overlaps_ignore, overlaps_ignore_indices, 1)
# gt max and indices, ignore max and indices
max_overlaps_normal = overlaps_normal[:, :top_k].flatten()
gt_assignment_normal = overlaps_normal_indices[:, :top_k].flatten()
max_overlaps_ignore = overlaps_ignore[:, :top_k].flatten()
gt_assignment_ignore = overlaps_ignore_indices[:, :top_k].flatten()
# cons masks
ignore_assign_mask = (max_overlaps_normal < fg_threshold).astype(
np.float32) * (max_overlaps_ignore > max_overlaps_normal).astype(
np.float32)
max_overlaps = max_overlaps_normal * (1 - ignore_assign_mask).astype(np.float32) + \
max_overlaps_ignore * ignore_assign_mask
gt_assignment = gt_assignment_normal * (1- ignore_assign_mask) + \
gt_assignment_ignore * ignore_assign_mask
gt_assignment = gt_assignment.astype(np.int32)
labels = gt_boxes_perimg[gt_assignment, 4]
fg_mask = (max_overlaps >= fg_threshold).astype(
np.float32) * (1 - F.equal(labels, config.ignore_label))
bg_mask = (max_overlaps < config.bg_threshold_high).astype(
np.float32) * (max_overlaps >= config.bg_threshold_low).astype(
np.float32)
fg_mask = fg_mask.reshape(-1, top_k)
bg_mask = bg_mask.reshape(-1, top_k)
pos_max = config.num_rois * config.fg_ratio
fg_inds_mask = _bernoulli_sample_masks(
fg_mask[:,
0], pos_max, 1) if sampling else F.equal(fg_mask[:, 0], 0)
neg_max = config.num_rois - fg_inds_mask.sum()
bg_inds_mask = _bernoulli_sample_masks(
bg_mask[:,
0], neg_max, 1) if sampling else F.equal(bg_mask[:, 0], 0)
labels = labels * fg_mask.reshape(-1)
keep_mask = fg_inds_mask + bg_inds_mask
keep_mask = keep_mask + F.equal(keep_mask.sum(), 0)
# keep_inds = mask_to_inds(keep_mask)
_, keep_inds = F.cond_take(keep_mask > 0, keep_mask)
#keep_inds = keep_inds[:F.minimum(config.num_rois, keep_inds.shapeof()[0])]
# labels
labels = labels.reshape(-1, top_k)[keep_inds]
gt_assignment = gt_assignment.reshape(
-1, top_k)[keep_inds].reshape(-1).astype(np.int32)
target_boxes = gt_boxes_perimg[gt_assignment, :4]
# rois = all_rois.ai[keep_inds]
rois = all_rois[keep_inds]
# target_shape = (rois.shapeof()[0], top_k, rois.shapeof()[-1])
n, c = rois.shape[0], rois.shape[1]
target_rois = F.broadcast_to(F.expand_dims(rois, 1),
(n, top_k, c)).reshape(-1, c)
# target_rois = F.add_axis(rois, 1).broadcast(target_shape).reshape(-1, rois.shapeof()[-1])
bbox_targets = bbox_transform_opr(target_rois[:, 1:5],
target_boxes[:, :4])
if config.rcnn_bbox_normalize_targets:
std_opr = mge.tensor(config.bbox_normalize_stds[None, :]).to(
rois.device)
mean_opr = mge.tensor(config.bbox_normalize_means[None, :]).to(
rois.device)
minus_opr = mean_opr / std_opr
bbox_targets = bbox_targets / std_opr - minus_opr
bbox_targets = bbox_targets.reshape(-1, top_k * 4)
return_rois.append(rois)
return_labels.append(labels)
return_bbox_targets.append(bbox_targets)
if config.batch_per_gpu == 1:
rois, labels, bbox_targets = rois.detach(), labels.detach(
), bbox_targets.detach()
return rois, labels, bbox_targets
# return F.zero_grad(rois), F.zero_grad(labels), F.zero_grad(bbox_targets)
else:
return_rois = F.concat(return_rois, axis=0)
return_labels = F.concat(return_labels, axis=0)
return_bbox_targets = F.concat(return_bbox_targets, axis=0)
return_rois = return_rois.detach()
return_labels = return_labels.detach()
return_bbox_targets = return_bbox_targets.detach()
return return_rois, return_labels, return_bbox_targets
def _anchor_double_target(gt_boxes, im_info, all_anchors):
gt_boxes, im_info = gt_boxes.detach(), im_info.detach()
all_anchors = all_anchors.detach()
gt_boxes = gt_boxes[:im_info[5].astype(np.int32), :]
dummy = -F.ones([1, gt_boxes.shape[1]]).to(gt_boxes.device)
gt_boxes = F.concat([gt_boxes, dummy], axis=0)
valid_mask = 1 - (gt_boxes[:, 4] < 0).astype(np.float32)
anchor_centers = _compute_center(all_anchors)
gtboxes_centers = _compute_center(gt_boxes)
# gtboxes_centers = gtboxes_centers * valid_mask.unsqueeze(1)
gtboxes_centers = gtboxes_centers * F.expand_dims(valid_mask, axis=1)
N, K = all_anchors.shape[0], gt_boxes.shape[0]
an_centers = F.expand_dims(anchor_centers, axis=1)
gt_centers = F.expand_dims(gtboxes_centers, axis=0)
# an_centers = anchor_centers.unsqueeze(1).repeat(1, K, 1)
# gt_centers = gtboxes_centers.unsqueeze(0).repeat(N, 1, 1)
distance = F.abs(an_centers - gt_centers)
distance = F.sqrt(F.pow(distance, 2).sum(axis=2))
start = 0
end = 5
overlaps = box_overlap_opr(all_anchors[:, :4], gt_boxes[:, :4])
overlaps *= F.expand_dims(valid_mask, axis=0)
default_num = 16
ious_list = []
for l in range(start, end):
_, index = F.cond_take(all_anchors[:, 4] == l, all_anchors[:, 4])
level_dist = distance[index, :].transpose(1, 0)
ious = overlaps[index, :].transpose(1, 0)
sorted_index = F.argsort(level_dist, descending=False)
n = min(sorted_index.shape[1], default_num)
ious = F.gather(ious, 1, sorted_index[:, :n]).transpose(1, 0)
ious_list.append(ious)
ious = F.concat(ious_list, axis=0)
mean_var = F.mean(ious, axis=0)
std_var = F.std(ious, 0)
iou_thresh_per_gt = mean_var + std_var
iou_thresh_per_gt = F.maximum(iou_thresh_per_gt, 0.2)
# limits the anchor centers in the gtboxes
N, K = all_anchors.shape[0], gt_boxes.shape[0]
anchor_points = an_centers
pos_area = _compute_pos_area(gt_boxes, 0.3)
# pos_area = pos_area.unsqueeze(0).repeat(N, 1, 1)
pos_area = F.broadcast_to(F.expand_dims(pos_area, axis=0),
(N, K, pos_area.shape[-1]))
l = anchor_points[:, :, 0] - pos_area[:, :, 0]
r = pos_area[:, :, 2] - anchor_points[:, :, 0]
t = anchor_points[:, :, 1] - pos_area[:, :, 1]
b = pos_area[:, :, 3] - anchor_points[:, :, 1]
is_in_gt = F.stack([l, r, t, b], axis=2)
is_in_gt = is_in_gt.min(axis=2) > 0.1
valid_mask = (overlaps >= F.expand_dims(
iou_thresh_per_gt, axis=0)) * is_in_gt.astype(np.float32)
ious = overlaps * valid_mask
sorted_index = F.argsort(ious, 1)
sorted_overlaps = F.gather(ious, 1, sorted_index)
max_overlaps = sorted_overlaps[:, :2].flatten()
argmax_overlaps = sorted_index[:, :2].flatten()
n, c = all_anchors.shape
device = all_anchors.device
labels = -F.ones(2 * n).to(device)
positive_mask = (max_overlaps >= 0.2).to(device).astype(np.float32)
negative_mask = (max_overlaps < 0.2).to(device).astype(np.float32)
labels = positive_mask + labels * (1 - positive_mask) * (1 - negative_mask)
bbox_targets = gt_boxes[argmax_overlaps, :4]
all_anchors = F.broadcast_to(F.expand_dims(all_anchors, axis=1),
(n, 2, c)).reshape(-1, c)
bbox_targets = bbox_transform_opr(all_anchors[:, :4], bbox_targets)
labels_cat = gt_boxes[argmax_overlaps, 4]
labels_cat = labels_cat * (1 - F.equal(labels, -1).astype(
np.float32)) - F.equal(labels, -1).astype(np.float32)
return labels, bbox_targets, labels_cat