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