def multiclass_nms_kpt(multi_bboxes,
multi_kpts,
multi_scores,
score_thr,
nms_cfg,
max_num=-1,
score_factors=None):
"""NMS for multi-class bboxes.
Args:
multi_bboxes (Tensor): shape (n, #class*4) or (n, 4)
multi_kpts (Tensor): shape (n, 34)
multi_scores (Tensor): shape (n, #class), where the 0th column
contains scores of the background class, but this will be ignored.
score_thr (float): bbox threshold, bboxes with scores lower than it
will not be considered.
nms_thr (float): NMS IoU threshold
max_num (int): if there are more than max_num bboxes after NMS,
only top max_num will be kept.
score_factors (Tensor): The factors multiplied to scores before
applying NMS
Returns:
tuple: (bboxes, labels), tensors of shape (k, 5) and (k, 1). Labels
are 0-based.
"""
num_classes = multi_scores.size(1) - 1
# exclude background category
if multi_bboxes.shape[1] > 4:
bboxes = multi_bboxes.view(multi_scores.size(0), -1, 4)
else:
bboxes = multi_bboxes[:, None].expand(-1, num_classes, 4)
scores = multi_scores[:, :-1]
# filter out boxes with low scores
valid_mask = scores > score_thr
bboxes = bboxes[valid_mask]
if score_factors is not None:
scores = scores * score_factors[:, None]
scores = scores[valid_mask]
labels = valid_mask.nonzero()[:, 1]
if bboxes.numel() == 0:
bboxes = multi_bboxes.new_zeros((0, 5))
labels = multi_bboxes.new_zeros((0, ), dtype=torch.long)
kpts = multi_bboxes.new_zeros((0, 35))
return bboxes, kpts, labels
if multi_kpts.shape[1] == 34:
kpts = multi_kpts[valid_mask, :]
else:
raise NotImplementedError
dets, keep = batched_nms(bboxes, scores, labels, nms_cfg)
if max_num > 0:
dets = dets[:max_num]
keep = keep[:max_num]
return dets, kpts[keep], labels[keep]
def _bboxes_nms(self, cls_scores, bboxes, score_factor, cfg):
max_scores, labels = torch.max(cls_scores, 1)
valid_mask = score_factor * max_scores >= cfg.score_thr
bboxes = bboxes[valid_mask]
scores = max_scores[valid_mask] * score_factor[valid_mask]
labels = labels[valid_mask]
if labels.numel() == 0:
return bboxes, labels
else:
dets, keep = batched_nms(bboxes, scores, labels, cfg.nms)
return dets, labels[keep]
def multiclass_nms(multi_bboxes,
multi_scores,
score_thr,
nms_cfg,
max_num=-1,
score_factors=None,
return_inds=False):
"""NMS for multi-class bboxes.
Args:
multi_bboxes (Tensor): shape (n, #class*4) or (n, 4)
multi_scores (Tensor): shape (n, #class), where the last column
contains scores of the background class, but this will be ignored.
score_thr (float): bbox threshold, bboxes with scores lower than it
will not be considered.
nms_thr (float): NMS IoU threshold
max_num (int, optional): if there are more than max_num bboxes after
NMS, only top max_num will be kept. Default to -1.
score_factors (Tensor, optional): The factors multiplied to scores
before applying NMS. Default to None.
return_inds (bool, optional): Whether return the indices of kept
bboxes. Default to False.
Returns:
tuple: (bboxes, labels, indices (optional)), tensors of shape (k, 5),
(k), and (k). Labels are 0-based.
"""
num_classes = multi_scores.size(1) - 1
# exclude background category
if multi_bboxes.shape[1] > 4:
bboxes = multi_bboxes.view(multi_scores.size(0), -1, 4)
else:
bboxes = multi_bboxes[:, None].expand(
multi_scores.size(0), num_classes, 4)
scores = multi_scores[:, :-1]
labels = torch.arange(num_classes, dtype=torch.long)
labels = labels.view(1, -1).expand_as(scores)
bboxes = bboxes.reshape(-1, 4)
scores = scores.reshape(-1)
labels = labels.reshape(-1)
if not torch.onnx.is_in_onnx_export():
# NonZero not supported in TensorRT
# remove low scoring boxes
valid_mask = scores > score_thr
# multiply score_factor after threshold to preserve more bboxes, improve
# mAP by 1% for YOLOv3
if score_factors is not None:
# expand the shape to match original shape of score
score_factors = score_factors.view(-1, 1).expand(
multi_scores.size(0), num_classes)
score_factors = score_factors.reshape(-1)
scores = scores * score_factors
if not torch.onnx.is_in_onnx_export():
# NonZero not supported in TensorRT
inds = valid_mask.nonzero(as_tuple=False).squeeze(1)
bboxes, scores, labels = bboxes[inds], scores[inds], labels[inds]
else:
# TensorRT NMS plugin has invalid output filled with -1
# add dummy data to make detection output correct.
bboxes = torch.cat([bboxes, bboxes.new_zeros(1, 4)], dim=0)
scores = torch.cat([scores, scores.new_zeros(1)], dim=0)
labels = torch.cat([labels, labels.new_zeros(1)], dim=0)
if bboxes.numel() == 0:
if torch.onnx.is_in_onnx_export():
raise RuntimeError('[ONNX Error] Can not record NMS '
'as it has not been executed this time')
if return_inds:
return bboxes, labels, inds
else:
return bboxes, labels
dets, keep = batched_nms(bboxes, scores, labels, nms_cfg)
if max_num > 0:
dets = dets[:max_num]
keep = keep[:max_num]
if return_inds:
return dets, labels[keep], keep
else:
return dets, labels[keep]
def multiclass_nms(multi_bboxes,
multi_scores,
score_thr,
nms_cfg,
max_num=-1,
score_factors=None):
"""NMS for multi-class bboxes.
Args:
multi_bboxes (Tensor): shape (n, #class*4) or (n, 4)
multi_scores (Tensor): shape (n, #class), where the last column
contains scores of the background class, but this will be ignored.
score_thr (float): bbox threshold, bboxes with scores lower than it
will not be considered.
nms_thr (float): NMS IoU threshold
max_num (int): if there are more than max_num bboxes after NMS,
only top max_num will be kept.
score_factors (Tensor): The factors multiplied to scores before
applying NMS
Returns:
tuple: (bboxes, labels), tensors of shape (k, 5) and (k, 1). Labels \
are 0-based.
"""
num_classes = multi_scores.size(1) - 1
# exclude background category
if multi_bboxes.shape[1] > 4:
bboxes = multi_bboxes.view(multi_scores.size(0), -1, 4)
else:
bboxes = multi_bboxes[:, None].expand(
multi_scores.size(0), num_classes, 4)
scores = multi_scores[:, :-1]
if score_factors is not None:
scores = scores * score_factors[:, None]
labels = torch.arange(num_classes, dtype=torch.long)
labels = labels.view(1, -1).expand_as(scores)
bboxes = bboxes.reshape(-1, 4)
scores = scores.reshape(-1)
labels = labels.reshape(-1)
# remove low scoring boxes
valid_mask = scores > score_thr
inds = valid_mask.nonzero(as_tuple=False).squeeze(1)
bboxes, scores, labels = bboxes[inds], scores[inds], labels[inds]
if inds.numel() == 0:
if torch.onnx.is_in_onnx_export():
raise RuntimeError('[ONNX Error] Can not record NMS '
'as it has not been executed this time')
return bboxes, labels
# TODO: add size check before feed into batched_nms
dets, keep = batched_nms(bboxes, scores, labels, nms_cfg)
if max_num > 0:
dets = dets[:max_num]
keep = keep[:max_num]
return dets, labels[keep]
def multiclass_nms_single(self, obj_scores, sem_scores, bbox, points,
input_meta):
"""Multi-class nms in single batch.
Args:
obj_scores (torch.Tensor): Objectness score of bounding boxes.
sem_scores (torch.Tensor): semantic class score of bounding boxes.
bbox (torch.Tensor): Predicted bounding boxes.
points (torch.Tensor): Input points.
input_meta (dict): Point cloud and image's meta info.
Returns:
tuple[torch.Tensor]: Bounding boxes, scores and labels.
"""
num_bbox = bbox.shape[0]
bbox = input_meta['box_type_3d'](
bbox.clone(),
box_dim=bbox.shape[-1],
with_yaw=self.bbox_coder.with_rot,
origin=(0.5, 0.5, 1.0))
if isinstance(bbox, LiDARInstance3DBoxes):
box_idx = bbox.points_in_boxes(points)
box_indices = box_idx.new_zeros([num_bbox + 1])
box_idx[box_idx == -1] = num_bbox
box_indices.scatter_add_(0, box_idx.long(),
box_idx.new_ones(box_idx.shape))
box_indices = box_indices[:-1]
nonempty_box_mask = box_indices >= 0
elif isinstance(bbox, DepthInstance3DBoxes):
box_indices = bbox.points_in_boxes(points)
nonempty_box_mask = box_indices.T.sum(1) >= 0
else:
raise NotImplementedError('Unsupported bbox type!')
corner3d = bbox.corners
minmax_box3d = corner3d.new(torch.Size((corner3d.shape[0], 6)))
minmax_box3d[:, :3] = torch.min(corner3d, dim=1)[0]
minmax_box3d[:, 3:] = torch.max(corner3d, dim=1)[0]
bbox_classes = torch.argmax(sem_scores, -1)
nms_selected = batched_nms(
minmax_box3d[nonempty_box_mask][:, [0, 1, 3, 4]],
obj_scores[nonempty_box_mask], bbox_classes[nonempty_box_mask],
self.test_cfg.nms_cfg)[1]
if nms_selected.shape[0] > self.test_cfg.max_output_num:
nms_selected = nms_selected[:self.test_cfg.max_output_num]
# filter empty boxes and boxes with low score
scores_mask = (obj_scores >= self.test_cfg.score_thr)
nonempty_box_inds = torch.nonzero(nonempty_box_mask).flatten()
nonempty_mask = torch.zeros_like(bbox_classes).scatter(
0, nonempty_box_inds[nms_selected], 1)
selected = (nonempty_mask.bool() & scores_mask.bool())
if self.test_cfg.per_class_proposal:
bbox_selected, score_selected, labels = [], [], []
for k in range(sem_scores.shape[-1]):
bbox_selected.append(bbox[selected].tensor)
score_selected.append(obj_scores[selected])
labels.append(
torch.zeros_like(bbox_classes[selected]).fill_(k))
bbox_selected = torch.cat(bbox_selected, 0)
score_selected = torch.cat(score_selected, 0)
labels = torch.cat(labels, 0)
else:
bbox_selected = bbox[selected].tensor
score_selected = obj_scores[selected]
labels = bbox_classes[selected]
return bbox_selected, score_selected, labels
def multiclass_nms(multi_bboxes,
multi_scores,
score_thr,
nms_cfg,
max_num=-1,
score_factors=None):
"""NMS for multi-class bboxes.
Args:
multi_bboxes (Tensor): shape (n, #class*4) or (n, 4)
multi_scores (Tensor): shape (n, #class), where the last column
contains scores of the background class, but this will be ignored.
score_thr (float): bbox threshold, bboxes with scores lower than it
will not be considered.
nms_thr (float): NMS IoU threshold
max_num (int): if there are more than max_num bboxes after NMS,
only top max_num will be kept.
score_factors (Tensor): The factors multiplied to scores before
applying NMS
Returns:
tuple: (bboxes, labels), tensors of shape (k, 5) and (k, 1). Labels \
are 0-based.
"""
num_classes = multi_scores.size(1) - 1
# exclude background category
if multi_bboxes.shape[1] > 4:
bboxes = multi_bboxes.view(multi_scores.size(0), -1, 4)
else:
bboxes = multi_bboxes[:, None].expand(multi_scores.size(0),
num_classes, 4)
scores = multi_scores[:, :-1]
# filter out boxes with low scores
valid_mask = scores > score_thr
# We use masked_select for ONNX exporting purpose,
# which is equivalent to bboxes = bboxes[valid_mask]
# (TODO): as ONNX does not support repeat now,
# we have to use this ugly code
bboxes = torch.masked_select(
bboxes,
torch.stack((valid_mask, valid_mask, valid_mask, valid_mask),
-1)).view(-1, 4)
if score_factors is not None:
scores = (scores * score_factors[:, None])
scores = torch.masked_select(scores, valid_mask)
labels = valid_mask.nonzero(as_tuple=False)[:, 1]
if bboxes.numel() == 0:
bboxes = multi_bboxes.new_zeros((0, 5))
labels = multi_bboxes.new_zeros((0, ), dtype=torch.long)
if torch.onnx.is_in_onnx_export():
raise RuntimeError('[ONNX Error] Can not record NMS '
'as it has not been executed this time')
return bboxes, labels
dets, keep = batched_nms(bboxes, scores, labels, nms_cfg)
if max_num > 0:
dets = dets[:max_num]
keep = keep[:max_num]
return dets, labels[keep]
def multiclass_nms(multi_bboxes,
multi_scores,
score_thr,
nms_cfg,
max_num=-1,
score_factors=None):
"""NMS for multi-class bboxes.
Args:
multi_bboxes (Tensor): shape (n, #class*4) or (n, 4)
multi_scores (Tensor): shape (n, #class), where the last column
contains scores of the background class, but this will be ignored.
score_thr (float): bbox threshold, bboxes with scores lower than it
will not be considered.
nms_thr (float): NMS IoU threshold
max_num (int): if there are more than max_num bboxes after NMS,
only top max_num will be kept.
score_factors (Tensor): The factors multiplied to scores before
applying NMS
Returns:
tuple: (bboxes, labels), tensors of shape (k, 5) and (k, 1). Labels \
are 0-based.
"""
iou_thr = nms_cfg['iou_threshold']
num_classes = multi_scores.size(1) - 1
# exclude background category
if multi_bboxes.shape[1] > 4:
bboxes = multi_bboxes.view(multi_scores.size(0), -1, 4)
else:
bboxes = multi_bboxes[:, None].expand(
multi_scores.size(0), num_classes, 4)
scores = multi_scores[:, :-1]
# filter out boxes with low scores
valid_mask = scores > score_thr
# We use masked_select for ONNX exporting purpose,
# which is equivalent to bboxes = bboxes[valid_mask]
# (TODO): as ONNX does not support repeat now,
# we have to use this ugly code
bboxes = torch.masked_select(
bboxes,
torch.stack((valid_mask, valid_mask, valid_mask, valid_mask),
-1)).view(-1, 4)
if score_factors is not None:
scores = scores * score_factors[:, None]
scores = torch.masked_select(scores, valid_mask)
labels = valid_mask.nonzero(as_tuple=False)[:, 1]
if bboxes.numel() == 0:
bboxes = multi_bboxes.new_zeros((0, 5))
labels = multi_bboxes.new_zeros((0, ), dtype=torch.long)
if torch.onnx.is_in_onnx_export():
raise RuntimeError('[ONNX Error] Can not record NMS '
'as it has not been executed this time')
return bboxes, labels
if nms_cfg['type']=='voting_cluster_diounms': # Score-voting Cluster-DIoU-NMS
scores, idx = scores.sort(0, descending=True)
bboxes = bboxes[idx]
labels = labels[idx]
box = bboxes + labels.unsqueeze(1).expand_as(bboxes)*4000
iouu = diou(box, box, 0.8)
iou = (iouu+0).triu_(diagonal=1)
B = iou
for i in range(999):
A=B
maxA = A.max(dim=0)[0]
E = (maxA <= iou_thr).float().unsqueeze(1).expand_as(A)
B=iou.mul(E)
if A.equal(B)==True:
break
# Now just filter out the ones higher than the threshold
B=torch.triu(iouu).mul(E)
keep = (maxA <= iou_thr)
weights = (torch.exp(-(1-(B*(B>0.7).float()))**2 / 0.025)) * (scores.reshape((1,len(scores))))
bboxes = torch.mm(weights, bboxes).float() / weights.sum(1, keepdim=True)
# Only keep the top max_num highest scores across all classes
if max_num > 0:
scores = scores[keep][:max_num]
labels = labels[keep][:max_num]
bboxes = bboxes[keep][:max_num]
dets = torch.cat([bboxes, scores[:, None]], dim=1)
elif nms_cfg['type']=='nms': # Original NMS
dets, keep = batched_nms(bboxes, scores, labels, nms_cfg)
if max_num > 0:
dets = dets[:max_num]
labels = labels[keep][:max_num]
else:
print("Error: The NMS function is unknown. Please check your cfg file.")
return dets, labels
def aug_test(self,
aug_x,
aug_proposal_boxes,
aug_proposal_features,
aug_img_metas,
aug_imgs_whwh,
rescale=False):
samples_per_gpu = len(aug_img_metas[0])
aug_det_bboxes = [[] for _ in range(samples_per_gpu)]
aug_det_labels = [[] for _ in range(samples_per_gpu)]
aug_mask_preds = [[] for _ in range(samples_per_gpu)]
for x, proposal_boxes, proposal_features, img_metas, imgs_whwh in \
zip(aug_x, aug_proposal_boxes, aug_proposal_features, aug_img_metas, aug_imgs_whwh):
num_imgs = len(img_metas)
proposal_list = [proposal_boxes[i] for i in range(num_imgs)]
ori_shapes = tuple(meta['ori_shape'] for meta in img_metas)
scale_factors = tuple(meta['scale_factor'] for meta in img_metas)
object_feats = proposal_features
for stage in range(self.num_stages):
rois = bbox2roi(proposal_list)
bbox_results = self._bbox_forward(stage, x, rois, object_feats,
img_metas)
object_feats = bbox_results['object_feats']
cls_score = bbox_results['cls_score']
proposal_list = bbox_results['detach_proposal_list']
if self.with_mask:
rois = bbox2roi(proposal_list)
mask_results = self._mask_forward(stage, x, rois,
bbox_results['attn_feats'])
mask_results['mask_pred'] = mask_results['mask_pred'].reshape(
num_imgs, -1, *mask_results['mask_pred'].size()[1:])
num_classes = self.bbox_head[-1].num_classes
det_bboxes = []
det_labels = []
if self.bbox_head[-1].loss_cls.use_sigmoid:
cls_score = cls_score.sigmoid()
else:
cls_score = cls_score.softmax(-1)[..., :-1]
for img_id in range(num_imgs):
cls_score_per_img = cls_score[img_id]
scores_per_img, topk_indices = cls_score_per_img.flatten(
0, 1).topk(self.test_cfg.max_per_img, sorted=False)
labels_per_img = topk_indices % num_classes
bbox_pred_per_img = proposal_list[img_id][topk_indices //
num_classes]
if rescale:
scale_factor = img_metas[img_id]['scale_factor']
bbox_pred_per_img /= bbox_pred_per_img.new_tensor(
scale_factor)
aug_det_bboxes[img_id].append(
torch.cat([bbox_pred_per_img, scores_per_img[:, None]],
dim=1))
det_bboxes.append(
torch.cat([bbox_pred_per_img, scores_per_img[:, None]],
dim=1))
aug_det_labels[img_id].append(labels_per_img)
det_labels.append(labels_per_img)
if self.with_mask:
if rescale and not isinstance(scale_factors[0], float):
scale_factors = [
torch.from_numpy(scale_factor).to(det_bboxes[0].device)
for scale_factor in scale_factors
]
_bboxes = [
det_bboxes[i][:, :4] *
scale_factors[i] if rescale else det_bboxes[i][:, :4]
for i in range(len(det_bboxes))
]
mask_pred = mask_results['mask_pred']
for img_id in range(num_imgs):
mask_pred_per_img = mask_pred[img_id].flatten(
0, 1)[topk_indices]
mask_pred_per_img = mask_pred_per_img[:, None, ...].repeat(
1, num_classes, 1, 1)
segm_result = self.mask_head[-1].get_seg_masks(
mask_pred_per_img,
_bboxes[img_id],
det_labels[img_id],
self.test_cfg,
ori_shapes[img_id],
scale_factors[img_id],
rescale,
format=False)
aug_mask_preds[img_id].append(
segm_result.detach().cpu().numpy())
det_bboxes, det_labels, mask_preds = [], [], []
for img_id in range(samples_per_gpu):
for aug_id in range(len(aug_det_bboxes[img_id])):
img_meta = aug_img_metas[aug_id][img_id]
img_shape = img_meta['ori_shape']
flip = img_meta['flip']
flip_direction = img_meta['flip_direction']
aug_det_bboxes[img_id][aug_id][:, :-1] = bbox_flip(
aug_det_bboxes[img_id][aug_id][:, :-1], img_shape,
flip_direction
) if flip else aug_det_bboxes[img_id][aug_id][:, :-1]
if flip:
if flip_direction == 'horizontal':
aug_mask_preds[img_id][aug_id] = aug_mask_preds[
img_id][aug_id][:, :, ::-1]
else:
aug_mask_preds[img_id][aug_id] = aug_mask_preds[
img_id][aug_id][:, ::-1, :]
for img_id in range(samples_per_gpu):
det_bboxes_per_im = torch.cat(aug_det_bboxes[img_id])
det_labels_per_im = torch.cat(aug_det_labels[img_id])
mask_preds_per_im = np.concatenate(aug_mask_preds[img_id])
# TODO(vealocia): implement batched_nms here.
det_bboxes_per_im, keep_inds = batched_nms(
det_bboxes_per_im[:, :-1], det_bboxes_per_im[:,
-1].contiguous(),
det_labels_per_im, self.test_cfg.nms)
det_bboxes_per_im = det_bboxes_per_im[:self.test_cfg.max_per_img,
...]
det_labels_per_im = det_labels_per_im[keep_inds][:self.test_cfg.
max_per_img, ...]
mask_preds_per_im = mask_preds_per_im[
keep_inds.detach().cpu().numpy()][:self.test_cfg.max_per_img,
...]
det_bboxes.append(det_bboxes_per_im)
det_labels.append(det_labels_per_im)
mask_preds.append(mask_preds_per_im)
ms_bbox_result = {}
ms_segm_result = {}
num_classes = self.bbox_head[-1].num_classes
bbox_results = [
bbox2result(det_bboxes[i], det_labels[i], num_classes)
for i in range(samples_per_gpu)
]
ms_bbox_result['ensemble'] = bbox_results
mask_results = [
mask2results(mask_preds[i], det_labels[i], num_classes)
for i in range(samples_per_gpu)
]
ms_segm_result['ensemble'] = mask_results
if self.with_mask:
results = list(
zip(ms_bbox_result['ensemble'], ms_segm_result['ensemble']))
else:
results = ms_bbox_result['ensemble']
return results
def perclass_nms(multi_bboxes,
multi_scores,
score_thr,
nms_cfg,
max_num=-1,
score_factors=None):
"""
This function has same input and output with `multiclass_nms`, but perform nms per class
Use it to save memory.
It contains many redundant code with `multiclass_nms`
"""
num_classes = multi_scores.size(1) - 1
# exclude background category
if multi_bboxes.shape[1] > 4:
bboxes = multi_bboxes.view(multi_scores.size(0), -1, 4)
else:
bboxes = multi_bboxes[:, None].expand(-1, num_classes, 4)
scores = multi_scores[:, :-1]
# filter out boxes with low scores
valid_mask = scores > score_thr
bboxes = bboxes[valid_mask]
if score_factors is not None:
scores = scores * score_factors[:, None]
scores = scores[valid_mask]
labels = valid_mask.nonzero()[:, 1]
if bboxes.numel() == 0:
bboxes = multi_bboxes.new_zeros((0, 5))
labels = multi_bboxes.new_zeros((0, ), dtype=torch.long)
return bboxes, labels
all_dets = []
all_labels = []
# do nms per class
for cls in range(num_classes):
cls_inds = labels == cls
cls_bboxes = bboxes[cls_inds]
if cls_bboxes.size(0) == 0:
continue
cls_scores = scores[cls_inds]
cls_labels = labels[cls_inds]
dets, keep = batched_nms(cls_bboxes, cls_scores, cls_labels, nms_cfg)
all_dets.append(dets)
all_labels.append(cls_labels[keep])
# concate the results -> sort by score -> select top n
all_dets = torch.cat(all_dets, dim=0)
all_labels = torch.cat(all_labels, dim=0)
_, sorted_ind = torch.sort(all_dets[:, -1], descending=True)
dets = all_dets[sorted_ind]
labels = all_labels[sorted_ind]
if max_num > 0:
dets = dets[:max_num]
labels = labels[:max_num]
return dets, labels
