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]
