def test_merge_aug_bboxes_3d(): if not torch.cuda.is_available(): pytest.skip('test requires GPU and torch+cuda') img_meta_0 = dict( pcd_horizontal_flip=False, pcd_vertical_flip=True, pcd_scale_factor=1.0) img_meta_1 = dict( pcd_horizontal_flip=True, pcd_vertical_flip=False, pcd_scale_factor=1.0) img_meta_2 = dict( pcd_horizontal_flip=False, pcd_vertical_flip=False, pcd_scale_factor=0.5) img_metas = [[img_meta_0], [img_meta_1], [img_meta_2]] boxes_3d = DepthInstance3DBoxes( torch.tensor( [[1.0473, 4.1687, -1.2317, 2.3021, 1.8876, 1.9696, 1.6956], [2.5831, 4.8117, -1.2733, 0.5852, 0.8832, 0.9733, 1.6500], [-1.0864, 1.9045, -1.2000, 0.7128, 1.5631, 2.1045, 0.1022]], device='cuda')) labels_3d = torch.tensor([0, 7, 6]) scores_3d = torch.tensor([0.5, 1.0, 1.0]) aug_result = dict( boxes_3d=boxes_3d, labels_3d=labels_3d, scores_3d=scores_3d) aug_results = [aug_result, aug_result, aug_result] test_cfg = mmcv.ConfigDict( use_rotate_nms=True, nms_across_levels=False, nms_thr=0.01, score_thr=0.1, min_bbox_size=0, nms_pre=100, max_num=50) results = merge_aug_bboxes_3d(aug_results, img_metas, test_cfg) expected_boxes_3d = torch.tensor( [[-1.0864, -1.9045, -1.2000, 0.7128, 1.5631, 2.1045, -0.1022], [1.0864, 1.9045, -1.2000, 0.7128, 1.5631, 2.1045, 3.0394], [-2.1728, 3.8090, -2.4000, 1.4256, 3.1262, 4.2090, 0.1022], [2.5831, -4.8117, -1.2733, 0.5852, 0.8832, 0.9733, -1.6500], [-2.5831, 4.8117, -1.2733, 0.5852, 0.8832, 0.9733, 1.4916], [5.1662, 9.6234, -2.5466, 1.1704, 1.7664, 1.9466, 1.6500], [1.0473, -4.1687, -1.2317, 2.3021, 1.8876, 1.9696, -1.6956], [-1.0473, 4.1687, -1.2317, 2.3021, 1.8876, 1.9696, 1.4460], [2.0946, 8.3374, -2.4634, 4.6042, 3.7752, 3.9392, 1.6956]]) expected_scores_3d = torch.tensor([ 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 0.5000, 0.5000, 0.5000 ]) expected_labels_3d = torch.tensor([6, 6, 6, 7, 7, 7, 0, 0, 0]) assert torch.allclose(results['boxes_3d'].tensor, expected_boxes_3d) assert torch.allclose(results['scores_3d'], expected_scores_3d) assert torch.all(results['labels_3d'] == expected_labels_3d)
def aug_test_pts(self, feats, img_metas, rescale=False): """Test function of point cloud branch with augmentaiton.""" # only support aug_test for one sample aug_bboxes = [] for x, img_meta in zip(feats, img_metas): outs = self.pts_bbox_head(x) bbox_list = self.pts_bbox_head.get_bboxes( *outs, img_meta, rescale=rescale) bbox_list = [ dict(boxes_3d=bboxes, scores_3d=scores, labels_3d=labels) for bboxes, scores, labels in bbox_list ] aug_bboxes.append(bbox_list[0]) # after merging, bboxes will be rescaled to the original image size merged_bboxes = merge_aug_bboxes_3d(aug_bboxes, img_metas, self.pts_bbox_head.test_cfg) return merged_bboxes
def aug_test(self, points, img_metas, imgs=None, rescale=False): """Test with augmentation.""" points_cat = [torch.stack(pts) for pts in points] feats = self.extract_feats(points_cat, img_metas) # only support aug_test for one sample aug_bboxes = [] for x, pts_cat, img_meta in zip(feats, points_cat, img_metas): bbox_preds = self.bbox_head(x, self.test_cfg.sample_mod) bbox_list = self.bbox_head.get_bboxes( pts_cat, bbox_preds, img_meta, rescale=rescale) bbox_list = [ dict(boxes_3d=bboxes, scores_3d=scores, labels_3d=labels) for bboxes, scores, labels in bbox_list ] aug_bboxes.append(bbox_list[0]) # after merging, bboxes will be rescaled to the original image size merged_bboxes = merge_aug_bboxes_3d(aug_bboxes, img_metas, self.bbox_head.test_cfg) return [merged_bboxes]
def aug_test(self, points, img_metas, imgs=None, rescale=False): """Test with augmentation.""" points_cat = [torch.stack(pts) for pts in points] feats_dict = self.extract_feats(points_cat, img_metas) for feat_dict in feats_dict: feat_dict['fp_xyz'] = [feat_dict['fp_xyz_net0'][-1]] feat_dict['fp_features'] = [feat_dict['hd_feature']] feat_dict['fp_indices'] = [feat_dict['fp_indices_net0'][-1]] # only support aug_test for one sample aug_bboxes = [] for feat_dict, pts_cat, img_meta in zip(feats_dict, points_cat, img_metas): if self.with_rpn: proposal_cfg = self.test_cfg.rpn rpn_outs = self.rpn_head(feat_dict, proposal_cfg.sample_mod) feat_dict.update(rpn_outs) # Generate rpn proposals proposal_list = self.rpn_head.get_bboxes( points, rpn_outs, img_metas, use_nms=proposal_cfg.use_nms) feat_dict['proposal_list'] = proposal_list else: raise NotImplementedError bbox_results = self.roi_head.simple_test( feat_dict, self.test_cfg.rcnn.sample_mod, img_meta, pts_cat, rescale=rescale) aug_bboxes.append(bbox_results) # after merging, bboxes will be rescaled to the original image size merged_bboxes = merge_aug_bboxes_3d(aug_bboxes, img_metas, self.bbox_head.test_cfg) return [merged_bboxes]
def aug_test_pts(self, feats, img_metas, rescale=False): """Test function of point cloud branch with augmentaiton. The function implementation process is as follows: - step 1: map features back for double-flip augmentation. - step 2: merge all features and generate boxes. - step 3: map boxes back for scale augmentation. - step 4: merge results. Args: feats (list[torch.Tensor]): Feature of point cloud. img_metas (list[dict]): Meta information of samples. rescale (bool): Whether to rescale bboxes. Default: False. Returns: dict: Returned bboxes consists of the following keys: - boxes_3d (:obj:`LiDARInstance3DBoxes`): Predicted bboxes. - scores_3d (torch.Tensor): Scores of predicted boxes. - labels_3d (torch.Tensor): Labels of predicted boxes. """ # only support aug_test for one sample outs_list = [] for x, img_meta in zip(feats, img_metas): outs = self.pts_bbox_head(x) # merge augmented outputs before decoding bboxes for task_id, out in enumerate(outs): for key in out[0].keys(): if img_meta[0]['pcd_horizontal_flip']: outs[task_id][0][key] = torch.flip( outs[task_id][0][key], dims=[2]) if key == 'reg': outs[task_id][0][key][:, 1, ...] = 1 - outs[ task_id][0][key][:, 1, ...] elif key == 'rot': outs[task_id][0][ key][:, 1, ...] = -outs[task_id][0][key][:, 1, ...] elif key == 'vel': outs[task_id][0][ key][:, 1, ...] = -outs[task_id][0][key][:, 1, ...] if img_meta[0]['pcd_vertical_flip']: outs[task_id][0][key] = torch.flip( outs[task_id][0][key], dims=[3]) if key == 'reg': outs[task_id][0][key][:, 0, ...] = 1 - outs[ task_id][0][key][:, 0, ...] elif key == 'rot': outs[task_id][0][ key][:, 0, ...] = -outs[task_id][0][key][:, 0, ...] elif key == 'vel': outs[task_id][0][ key][:, 0, ...] = -outs[task_id][0][key][:, 0, ...] outs_list.append(outs) preds_dicts = dict() scale_img_metas = [] # concat outputs sharing the same pcd_scale_factor for i, (img_meta, outs) in enumerate(zip(img_metas, outs_list)): pcd_scale_factor = img_meta[0]['pcd_scale_factor'] if pcd_scale_factor not in preds_dicts.keys(): preds_dicts[pcd_scale_factor] = outs scale_img_metas.append(img_meta) else: for task_id, out in enumerate(outs): for key in out[0].keys(): preds_dicts[pcd_scale_factor][task_id][0][key] += out[ 0][key] aug_bboxes = [] for pcd_scale_factor, preds_dict in preds_dicts.items(): for task_id, pred_dict in enumerate(preds_dict): # merge outputs with different flips before decoding bboxes for key in pred_dict[0].keys(): preds_dict[task_id][0][key] /= len(outs_list) / len( preds_dicts.keys()) bbox_list = self.pts_bbox_head.get_bboxes(preds_dict, img_metas[0], rescale=rescale) bbox_list = [ dict(boxes_3d=bboxes, scores_3d=scores, labels_3d=labels) for bboxes, scores, labels in bbox_list ] aug_bboxes.append(bbox_list[0]) if len(preds_dicts.keys()) > 1: # merge outputs with different scales after decoding bboxes merged_bboxes = merge_aug_bboxes_3d(aug_bboxes, scale_img_metas, self.pts_bbox_head.test_cfg) return merged_bboxes else: for key in bbox_list[0].keys(): bbox_list[0][key] = bbox_list[0][key].to('cpu') import pdb pdb.set_trace() return bbox_list[0]
def aug_test(self, points=None, img_metas=None, imgs=None, bboxes_2d=None, rescale=False, **kwargs): """Test function with augmentation, stage 2. Args: points (list[list[torch.Tensor]], optional): the outer list indicates test-time augmentations and the inner list contains all points in the batch, where each Tensor should have a shape NxC. Defaults to None. img_metas (list[list[dict]], optional): the outer list indicates test-time augs (multiscale, flip, etc.) and the inner list indicates images in a batch. Defaults to None. imgs (list[list[torch.Tensor]], optional): the outer list indicates test-time augmentations and inner Tensor should have a shape NxCxHxW, which contains all images in the batch. Defaults to None. Defaults to None. bboxes_2d (list[list[torch.Tensor]], optional): Provided 2d bboxes, not supported yet. Defaults to None. rescale (bool, optional): Whether or not rescale bboxes. Defaults to False. Returns: list[dict]: Predicted 3d boxes. """ points_cat = [torch.stack(pts) for pts in points] feats = self.extract_pts_feats(points_cat, img_metas) # only support aug_test for one sample aug_bboxes = [] for x, pts_cat, img_meta, bbox_2d, img in zip(feats, points_cat, img_metas, bboxes_2d, imgs): bbox_2d = self.extract_bboxes_2d( img, img_metas, train=False, bboxes_2d=bbox_2d, **kwargs) seeds_3d, seed_3d_features, seed_indices = x img_features, masks = self.fusion_layer(img, bbox_2d, seeds_3d, img_metas) inds = sample_valid_seeds(masks, self.num_sampled_seed) batch_size, img_feat_size = img_features.shape[:2] pts_feat_size = seed_3d_features.shape[1] inds_img = inds.view(batch_size, 1, -1).expand(-1, img_feat_size, -1) img_features = img_features.gather(-1, inds_img) inds = inds % inds.shape[1] inds_seed_xyz = inds.view(batch_size, -1, 1).expand(-1, -1, 3) seeds_3d = seeds_3d.gather(1, inds_seed_xyz) inds_seed_feats = inds.view(batch_size, 1, -1).expand(-1, pts_feat_size, -1) seed_3d_features = seed_3d_features.gather(-1, inds_seed_feats) seed_indices = seed_indices.gather(1, inds) img_features = self.img_mlp(img_features) fused_features = torch.cat([seed_3d_features, img_features], dim=1) feat_dict = dict( seed_points=seeds_3d, seed_features=fused_features, seed_indices=seed_indices) bbox_preds = self.pts_bbox_head_joint(feat_dict, self.test_cfg.pts.sample_mod) bbox_list = self.pts_bbox_head_joint.get_bboxes( pts_cat, bbox_preds, img_metas, rescale=rescale) bbox_list = [ dict(boxes_3d=bboxes, scores_3d=scores, labels_3d=labels) for bboxes, scores, labels in bbox_list ] aug_bboxes.append(bbox_list[0]) # after merging, bboxes will be rescaled to the original image size merged_bboxes = merge_aug_bboxes_3d(aug_bboxes, img_metas, self.bbox_head.test_cfg) return [merged_bboxes]