def forward(self, pred_cls_list, rpn_num_prob_list, pred_reg_list,
anchors_list, rpn_iou_list, boxes, im_info):
all_anchors_list = [
F.concat([a, i * F.ones([a.shape[0], 1]).to(a.device)], axis=1)
for i, a in enumerate(anchors_list)
]
all_anchors_final = F.concat(all_anchors_list, axis=0)
rpn_bbox_offset_final = F.concat(pred_reg_list, axis=1)
rpn_cls_prob_final = F.concat(pred_cls_list, axis=1)
rpn_iou_prob_final = F.concat(rpn_iou_list, axis=1)
rpn_num_per_points_final = F.concat(rpn_num_prob_list, axis=1)
rpn_labels, rpn_target_boxes = rpn_anchor_target_opr(
boxes, im_info, all_anchors_final)
ious_target = self.anchor_iou_target_opr(boxes, im_info,
all_anchors_final,
rpn_bbox_offset_final)
n = rpn_labels.shape[0]
target_boxes = rpn_target_boxes.reshape(n, -1, 4)
rpn_cls_prob_final = rpn_cls_prob_final.reshape(n, -1, 1)
offsets_final = rpn_bbox_offset_final.reshape(n, -1, 4)
rpn_labels = rpn_labels.transpose(2, 0, 1)
a, b = rpn_labels[0], rpn_labels[1]
ignores = b - F.equal(a, 0).astype(np.float32) * F.equal(b, 0).astype(
np.float32)
labels = F.stack([a, ignores], axis=2).reshape(n, -1)
cls_loss = sigmoid_cross_entropy_retina(rpn_cls_prob_final,
labels,
alpha=config.focal_loss_alpha,
gamma=config.focal_loss_gamma)
rpn_bbox_loss = smooth_l1_loss_retina(offsets_final, target_boxes,
labels)
rpn_labels = labels.reshape(n, -1, 2)
rpn_iou_loss = iou_l1_loss(rpn_iou_prob_final, ious_target, rpn_labels)
# whether one anchor produce one proposal or two.
nlabels = ((labels.reshape(n, -1, 2) > 0).sum(2)).flatten() - 1
c = rpn_num_per_points_final.shape[2]
num_per_anchor = rpn_num_per_points_final.reshape(-1, c)
rpn_num_per_points_final = rpn_num_per_points_final.reshape(-1, c)
nlabels = nlabels.reshape(-1)
rpn_num_loss = softmax_loss(rpn_num_per_points_final, nlabels)
loss_dict = {}
loss_dict['rpn_cls_loss'] = cls_loss
loss_dict['rpn_bbox_loss'] = 2 * rpn_bbox_loss
loss_dict['rpn_iou_loss'] = 2 * rpn_iou_loss
loss_dict['rpn_num_loss'] = rpn_num_loss
return loss_dict
def iou_l1_loss(pred, max_overlaps, gt, ignore_label=-1, background=0):
pred = pred.reshape(pred.shape[0], -1, max_overlaps.shape[2])
abs_x = F.abs(pred - max_overlaps)
mask_bg = 1 - F.equal(gt, background).astype(np.float32)
mask_ig = 1 - F.equal(gt, ignore_label).astype(np.float32)
mask = mask_bg * mask_ig
mask = mask.reshape(mask.shape[0], -1, pred.shape[2])
loss = (abs_x * mask).sum() / F.maximum(mask.sum(), 1)
return loss
def mask_anchor_opr(gtboxes, im_info, anchors, labels):
eps = 1e-6
gtboxes = gtboxes[:im_info[5].astype(np.int32), :]
ignore_mask = (gtboxes[:, 4] < 0).astype(np.float32)
mask_flag = F.zeros(labels.shape[0])
N, K = anchors.shape[0], gtboxes.shape[0]
p_pred = F.broadcast_to(F.expand_dims(anchors, 1),
(N, K, anchors.shape[1]))
p_gt = F.broadcast_to(F.expand_dims(gtboxes, 0), (N, K, gtboxes.shape[1]))
max_off = F.concat([
F.maximum(p_pred[:, :, :2], p_gt[:, :, :2]),
F.minimum(p_pred[:, :, 2:4], p_gt[:, :, 2:4])
],
axis=2)
I = F.maximum(max_off[:, :, 2] - max_off[:, :, 0] + 1, 0) * F.maximum(
max_off[:, :, 3] - max_off[:, :, 1] + 1, 0)
A = F.maximum(p_pred[:, :, 2] - p_pred[:, :, 0] + 1, 0) * F.maximum(
p_pred[:, :, 3] - p_pred[:, :, 1] + 1, 0)
# I = F.maximum(I, 0)
# A = F.maximum(A, 0)
IoA = I / (A + eps)
IoA = IoA * F.expand_dims(ignore_mask, 0)
mask_flag = (IoA > 0.5).sum(axis=1) > 0
labels = labels - F.equal(labels, 0).astype(np.float32) * mask_flag.astype(
np.float32)
return labels
def fpn_anchor_target(boxes, im_info, all_anchors_list):
final_labels_list = []
final_bbox_targets_list = []
for bid in range(config.batch_per_gpu):
batch_labels_list = []
batch_bbox_targets_list = []
for i in range(len(all_anchors_list)):
anchors_perlvl = all_anchors_list[i]
rpn_labels_perlvl, rpn_bbox_targets_perlvl = fpn_anchor_target_opr_core_impl(
boxes[bid], im_info[bid], anchors_perlvl)
batch_labels_list.append(rpn_labels_perlvl)
batch_bbox_targets_list.append(rpn_bbox_targets_perlvl)
# here we samples the rpn_labels
concated_batch_labels = F.concat(batch_labels_list, axis=0)
concated_batch_bbox_targets = F.concat(batch_bbox_targets_list, axis=0)
# sample labels
num_positive = config.num_sample_anchors * config.positive_anchor_ratio
concated_batch_labels = _bernoulli_sample_labels(
concated_batch_labels, num_positive, 1, config.ignore_label)
num_positive = F.equal(concated_batch_labels, 1).sum()
num_negative = config.num_sample_anchors - num_positive
concated_batch_labels = _bernoulli_sample_labels(
concated_batch_labels, num_negative, 0, config.ignore_label)
final_labels_list.append(concated_batch_labels)
final_bbox_targets_list.append(concated_batch_bbox_targets)
final_labels = F.concat(final_labels_list, axis=0)
final_bbox_targets = F.concat(final_bbox_targets_list, axis=0)
return F.zero_grad(final_labels), F.zero_grad(final_bbox_targets)
def anchor_iou_target_opr(self, boxes, im_info, all_anchors,
rpn_bbox_offsets):
n = rpn_bbox_offsets.shape[0]
res = []
for i in range(n):
gtboxes = boxes[i, :im_info[i, 5].astype(np.int32)]
offsets = rpn_bbox_offsets[i].reshape(-1, 4).detach()
m = offsets.shape[0]
an, ac = all_anchors.shape[0], all_anchors.shape[1]
anchors = F.broadcast_to(F.expand_dims(all_anchors, 1),
(an, 1, ac)).reshape(-1, ac)
dtboxes = bbox_transform_inv_opr(anchors[:, :4], offsets[:, :4])
overlaps = box_overlap_opr(dtboxes, gtboxes[:, :4])
ignore_mask = 1 - F.equal(
gtboxes[:, 4], config.anchor_ignore_label).astype(np.float32)
ignore_mask = F.expand_dims(ignore_mask, axis=0)
overlaps = overlaps * ignore_mask
index = F.argmax(overlaps, axis=1)
value = F.nn.indexing_one_hot(overlaps, index, 1)
value = F.expand_dims(F.expand_dims(value, axis=1), axis=0)
res.append(value)
result = F.concat(res, 0)
return result
def softmax_loss(pred, label, ignore_label=-1):
max_pred = F.zero_grad(pred.max(axis=1, keepdims=True))
pred -= max_pred
log_prob = pred - F.log(F.exp(pred).sum(axis=1, keepdims=True))
mask = 1 - F.equal(label, ignore_label)
vlabel = label * mask
loss = -(F.indexing_one_hot(log_prob, vlabel, 1) * mask)
return loss
def fpn_anchor_target_opr_core_impl(gt_boxes,
im_info,
anchors,
allow_low_quality_matches=True):
ignore_label = config.ignore_label
# get the gt boxes
gtboxes = gt_boxes[:im_info[5].astype(np.int32)]
ignore_mask = F.equal(gtboxes[:, 4], config.ignore_label)
# find the valid gtboxes
_, index = F.cond_take(1 - ignore_mask > 0, ignore_mask)
valid_gt_boxes = gtboxes[index.astype(np.int32)]
# compute the iou matrix
overlaps = box_overlap_opr(anchors, valid_gt_boxes[:, :4])
# match the dtboxes
a_shp0 = anchors.shape[0]
argmax_overlaps = F.argmax(overlaps, axis=1)
max_overlaps = F.nn.indexing_one_hot(overlaps,
argmax_overlaps.astype(np.int32), 1)
labels = F.ones(a_shp0).astype(np.int32) * ignore_label
# set negative ones
labels = labels * (max_overlaps >= config.rpn_negative_overlap).astype(
np.float32)
# set positive ones
fg_mask = (max_overlaps >= config.rpn_positive_overlap)
const_one = mge.tensor(1.0)
if allow_low_quality_matches:
# match the max gt
gt_max_overlaps = F.max(overlaps, axis=0)
gt_argmax_overlaps = F.argmax(overlaps, axis=0)
gt_argmax_overlaps = gt_argmax_overlaps.astype(np.int32)
max_overlaps[gt_argmax_overlaps] = 1.
m = gt_max_overlaps.shape[0]
argmax_overlaps[gt_argmax_overlaps] = F.linspace(0, m - 1,
m).astype(np.int32)
fg_mask = (max_overlaps >= config.rpn_positive_overlap)
labels[fg_mask] = 1
# compute the bbox targets
bbox_targets = bbox_transform_opr(anchors,
valid_gt_boxes[argmax_overlaps, :4])
if config.rpn_bbox_normalize_targets:
std_opr = mge.tensor(config.bbox_normalize_stds[None, :]).to(
anchors.device)
mean_opr = mge.tensor(config.bbox_normalize_means[None, :]).to(
anchors.device)
minus_opr = mean_opr / std_opr
bbox_targets = bbox_targets / std_opr - minus_opr
return labels, bbox_targets
def _bernoulli_sample_masks(masks, num_samples, sample_value):
""" Using the bernoulli sampling method"""
sample_mask = F.equal(masks, sample_value)
num_mask = sample_mask.sum()
num_final_samples = F.minimum(num_mask, num_samples)
# here, we use the bernoulli probability to sample the anchors
sample_prob = num_final_samples / num_mask
uniform_rng = rand.uniform(sample_mask.shapeof()[0])
after_sampled_mask = (uniform_rng <= sample_prob) * sample_mask
return after_sampled_mask
def softmax_cross_entropy(pred, label, axis=1, ignore_index=255):
offset = F.zero_grad(pred.max(axis=axis, keepdims=True))
pred = pred - offset
log_prob = pred - F.log(F.exp(pred).sum(axis=axis, keepdims=True))
mask = 1 - F.equal(label, ignore_index)
vlabel = label * mask
loss = -(F.indexing_one_hot(log_prob, vlabel, axis) *
mask).sum() / F.maximum(mask.sum(), 1)
return loss
def softmax_loss_opr(pred, label, ignore_label=-1):
max_pred = pred.max(axis=1, keepdims=True).detach()
pred -= max_pred
log_prob = pred - F.log(F.exp(pred).sum(axis=1, keepdims=True))
mask = 1 - F.equal(label, ignore_label)
vlabel = label * mask.astype(np.float32)
loss = -(F.nn.indexing_one_hot(log_prob, vlabel.astype(np.int32),
1).flatten() * mask)
return loss
def box_overlap_ignore_opr(box: Tensor, gt: Tensor, ignore_label=-1) -> Tensor:
"""
Given two lists of boxes of size N and M,
compute the IoU (intersection over union)
between __all__ N x M pairs of boxes.
The box order must be (xmin, ymin, xmax, ymax).
Args:
boxes1,boxes2 (Boxes): two `Boxes`. Contains N & M boxes, respectively.
Returns:
Tensor: IoU, sized [N,M].
"""
# box = boxes1
# gt = boxes2
# target_shape = (boxes1.shapeof()[0], boxes2.shapeof()[0], 4)
eps = 1e-5
N, K = box.shape[0], gt.shape[0]
b_box = F.broadcast_to(F.expand_dims(box, 1), (N, K, box.shape[1]))
b_gt = F.broadcast_to(F.expand_dims(gt, 0), (N, K, gt.shape[1]))
# b_box = F.add_axis(boxes1, 1).broadcast(*target_shape)
# b_gt = F.add_axis(boxes2[:, :4], 0).broadcast(*target_shape)
iw = F.minimum(b_box[:, :, 2], b_gt[:, :, 2]) - F.maximum(
b_box[:, :, 0], b_gt[:, :, 0])
ih = F.minimum(b_box[:, :, 3], b_gt[:, :, 3]) - F.maximum(
b_box[:, :, 1], b_gt[:, :, 1])
inter = F.maximum(iw, 0) * F.maximum(ih, 0)
area_box = F.maximum(box[:, 2] - box[:, 0], 0) * F.maximum(
box[:, 3] - box[:, 1], 0)
area_gt = F.maximum(gt[:, 2] - gt[:, 0], 0) * F.maximum(
gt[:, 3] - gt[:, 1], 0)
# area_target_shape = (box.shapeof()[0], gt.shapeof()[0])
# b_area_box = F.add_axis(area_box, 1).broadcast(*area_target_shape)
# b_area_gt = F.add_axis(area_gt, 0).broadcast(*area_target_shape)
b_area_box = F.broadcast_to(F.expand_dims(area_box, 1), (N, K)) + eps
b_area_gt = F.broadcast_to(F.expand_dims(area_gt, 0), (N, K))
union = b_area_box + b_area_gt - inter + eps
overlaps_normal = F.maximum(inter / union, 0)
overlaps_ignore = F.maximum(inter / b_area_box, 0)
overlaps = F.maximum(inter / union, 0)
# gt_ignore_mask = F.add_axis(F.equal(gt[:, 4], ignore_label), 0).broadcast(*area_target_shape)
ignore_mask = F.equal(gt[:, 4], ignore_label)
gt_ignore_mask = F.expand_dims(ignore_mask, 0)
overlaps_normal *= (1 - gt_ignore_mask)
overlaps_ignore *= gt_ignore_mask
return overlaps_normal, overlaps_ignore
def forward(self, fpn_fms, rcnn_rois, gt_boxes=None, im_info=None):
if self.training:
loss = {}
for i, _ in enumerate(self.iou_thrs):
loss_dict, prob = self.subnets[i](fpn_fms, rcnn_rois, gt_boxes,
im_info)
rois = prob[:, 1]
rcnn_list = []
for bid in range(config.batch_per_gpu):
mask = F.equal(rois[:, 5], bid)
_, index = F.cond_take(mask > 0, mask)
batch_id = bid * F.ones([mask.sum(), 1])
m = F.concat([batch_id, rois[index, :4]], axis=1)
rcnn_list.append(m)
rcnn_rois = F.concat(rcnn_list, axis=0)
loss.update(loss_dict)
return loss
else:
# boxes_pred = self._forward_test(fpn_fms, rcnn_rois)
for i, _ in enumerate(self.iou_thrs):
prob = self.subnets[i](fpn_fms, rcnn_rois)
rois = prob[:, 1]
rcnn_list = []
for bid in range(1):
mask = F.equal(rois[:, 5], bid)
_, index = F.cond_take(mask > 0, mask)
batch_id = bid * F.ones([mask.sum(), 1])
m = F.concat([batch_id, rois[index, :4]], axis=1)
rcnn_list.append(m)
rcnn_rois = F.concat(rcnn_list, axis=0)
return prob[:, :, :5]
def _bernoulli_sample_labels(labels,
num_samples,
sample_value,
ignore_label=-1):
""" Using the bernoulli sampling method"""
sample_label_mask = F.equal(labels, sample_value)
num_mask = sample_label_mask.sum()
num_final_samples = F.minimum(num_mask, num_samples)
# here, we use the bernoulli probability to sample the anchors
sample_prob = num_final_samples / num_mask
uniform_rng = rand.uniform(sample_label_mask.shapeof()[0])
disable_mask = (uniform_rng >= sample_prob) * sample_label_mask
#TODO check cudaerror: illegal memory access was encountered
labels = labels * (1 - disable_mask) + disable_mask * ignore_label
return labels
def rpn_anchor_target_opr(gt_boxes, im_info, anchors):
rpn_label_list, rpn_target_boxes_list, iou_thresh_list = [], [], []
for i in range(config.train_batch_per_gpu):
rpn_labels, rpn_target_boxes, _ = _anchor_double_target(
gt_boxes[i], im_info[i], anchors)
rpn_labels = rpn_labels.reshape(-1, 2)
c = rpn_target_boxes.shape[1]
rpn_target_boxes = rpn_target_boxes.reshape(-1, 2, c)
# mask the anchors overlapping with ignore regions
ignore_label = mask_anchor_opr(gt_boxes[i], im_info[i], anchors,
rpn_labels[:, 0])
rpn_labels = rpn_labels - F.equal(rpn_labels, 0).astype(
np.float32) * F.expand_dims(ignore_label < 0, 1).astype(np.float32)
# rpn_labels = rpn_labels - rpn_labels.eq(0).astype(np.float32) * (ignore_label < 0).unsqueeze(1).astype(np.float32)
rpn_label_list.append(F.expand_dims(rpn_labels, 0))
rpn_target_boxes_list.append(F.expand_dims(rpn_target_boxes, 0))
rpn_labels = F.concat(rpn_label_list, axis=0)
rpn_target_boxes = F.concat(rpn_target_boxes_list, axis=0)
return rpn_labels, rpn_target_boxes
def sigmoid_cross_entropy_retina(pred,
label,
ignore_label=-1,
background=0,
alpha=0.5,
gamma=0):
device = pred.device
mask = 1 - F.equal(label, ignore_label).astype(np.float32)
vlabel = label * mask
n, m, c = pred.shape
zero_mat = F.zeros([n, m, c + 1]).to(device)
index = F.expand_dims(vlabel, 2).astype(np.int32)
one_hot = F.scatter(zero_mat, 2, index, F.ones([n, m, 1]))
onehot = one_hot[:, :, 1:]
pos_part = F.pow(1 - pred, gamma) * onehot * F.log(pred)
neg_part = F.pow(pred, gamma) * (1 - onehot) * F.log(1 - pred)
loss = -(alpha * pos_part + (1 - alpha) * neg_part).sum(axis=2) * mask
positive_mask = (label > 0)
return loss.sum() / F.maximum(positive_mask.sum(), 1)
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
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 filter_boxes_opr(boxes, min_size):
"""Remove all boxes with any side smaller than min_size."""
wh = boxes[:, 2:4] - boxes[:, 0:2] + 1
keep_mask = F.prod(wh >= min_size, axis=1).astype(np.float32)
keep_mask = keep_mask + F.equal(keep_mask.sum(), 0).astype(np.float32)
return keep
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 forward(self, features, label=None, mask=None):
"""
if label and mask both None, the loss will degenerate to
SimSLR unsupervised loss.
Reference:
"A Simple Framework for Contrastive Learning of Visual Representations"<https://arxiv.org/pdf/2002.05709.pdf>
"Supervised Contrastive Learning"<https://arxiv.org/abs/2004.11362>
Args:
features(tensor): The embedding feature. shape=[bs, n_views, ...]
label(tensor): The label of images, shape=[bs]
mask(tensor): contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
has the same class as sample i. Can be asymmetric.
return:
loss
"""
if len(features.shape) < 3:
raise ValueError("Features need have 3 dimensions at least")
bs, num_view = features.shape[:2]
#if dimension > 3, change the shape of the features to [bs, num_view, ...]
if len(features.shape) > 3:
features = features.reshape(bs, num_view, -1)
#label and mask cannot provided at the same time
if (label is not None) and (mask is not None):
raise ValueError("label and mask cannot provided at the same time")
elif (label is None) and (mask is None):
mask = F.eye(bs, dtype="float32")
elif label is not None:
label = label.reshape(-1, 1)
if label.shape[0] != bs:
raise RuntimeError(
"Num of labels does not match num of features")
mask = F.equal(label, label.T)
else:
mask = mask.astype("float32")
contrast_count = features.shape[1]
features = F.split(features, features.shape[1], axis=1)
contrast_feature = F.squeeze(F.concat(features, axis=0), axis=1)
if self.contrast_mode == "one":
anchor_feature = features[:, 0]
anchor_count = 1
elif self.contrast_mode == "all":
anchor_feature = contrast_feature
anchor_count = contrast_count
else:
raise ValueError("Unknown mode:{}".format(self.contrast_mode))
#compute logits
anchor_dot_contrast = F.div(
F.matmul(anchor_feature, contrast_feature.T), self.temperate)
#for numerical stability
logits_max = F.max(anchor_dot_contrast, axis=-1, keepdims=True)
logits = anchor_dot_contrast - logits_max
#tile mask
an1, con = mask.shape[:2]
nums = anchor_count * contrast_count
# mask-out self-contrast cases
mask = F.stack([mask] * nums).reshape(an1 * anchor_count,
con * contrast_count)
logits_mask = F.scatter(
F.ones_like(mask), 1,
F.arange(0, int(bs * anchor_count), dtype="int32").reshape(-1, 1),
F.zeros(int(bs * anchor_count), dtype="int32").reshape(-1, 1))
mask = mask * logits_mask
#compute log_prob
exp_logits = F.exp(logits) * logits_mask
log_prob = logits - F.log(F.sum(exp_logits, axis=1,
keepdims=True)) #equation 2
#mean
mean_log_prob_pos = F.sum(mask * log_prob, axis=1) / F.sum(mask,
axis=1)
#loss
loss = -(self.temperate / self.base_temperate) * mean_log_prob_pos
loss = F.mean(loss.reshape(anchor_count, bs))
return loss
def _get_mask_of_label(label, background, ignore_label):
mask_fg = 1 - F.equal(label, background).astype(np.float32)
mask_ig = 1 - F.equal(label, ignore_label).astype(np.float32)
mask = mask_fg * mask_ig
return mask, mask_ig
def roi_pool(rpn_fms,
rois,
stride,
pool_shape,
roi_type='roi_align',
labels=None,
bbox_targets=None):
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_sizes = F.sqrt((rois[:, 3] - rois[:, 1]) * (rois[:, 4] - rois[:, 2]))
level_assignments = F.floor(canonical_level +
F.log(box_sizes / 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
available_masks = F.concat(
[F.ones(level_assignments.shape[0]),
F.zeros(num_fms)], axis=0)
level_assignments = F.concat(
[level_assignments,
mge.tensor(np.arange(num_fms, dtype=np.int32))],
axis=0)
rois = F.concat([rois, F.zeros((num_fms, rois.shape[-1]))], axis=0)
if labels is not None and bbox_targets is not None:
labels = F.concat([labels, F.ones((num_fms, labels.shape[-1]))],
axis=0)
bbox_targets = F.concat(
[bbox_targets,
F.zeros((num_fms, bbox_targets.shape[-1]))], axis=0)
pool_list, inds_list = [], []
for i in range(len(rpn_fms)):
# mask = level_assignments == i
# inds = mask_to_inds(mask)
mask = F.equal(level_assignments, i)
_, inds = F.cond_take(mask > 0, mask)
rois_fm = rois[inds.astype(np.int32)]
if roi_type == 'roi_pool':
pool_fm = F.nn.roi_pooling(rpn_fms[i],
rois_fm,
pool_shape,
mode='max',
scale=1.0 / stride[i])
elif roi_type == 'roi_align':
pool_fm = F.nn.roi_align(rpn_fms[i],
rois_fm,
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)
pool_feature = F.concat(pool_list, axis=0)
ordered_available_masks = available_masks[fm_order]
# available_inds = mask_to_inds(ordered_available_masks)
_, available_inds = F.cond_take(ordered_available_masks > 0,
ordered_available_masks)
available_inds = available_inds.astype(np.int32)
pool_feature = pool_feature[available_inds.astype(np.int32)]
rois = rois[fm_order, :][available_inds.astype(np.int32)]
if labels is not None:
labels = labels[fm_order][available_inds]
bbox_targets = bbox_targets[fm_order][available_inds]
return pool_feature, rois, labels.detach(), bbox_targets.detach()
else:
return pool_feature, rois, None, None
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
