[feature]: support ssd_3d_head in 3DSSD (#83)

* add ssd3dhead

* fix bugs for anchorfreebboxcoder

* modify ssd 3d head

* modify ssd 3d head

* reconstruct ssd3dhead and votehead

* add unittest

* modify 3dssd config

* modify 3dssd head

* modify 3dssd head

* rename base conv bbox head

* modify vote module

* modify 3dssd config

* fix bugs for unittest

* modify test_heads.py

* fix bugs for h3d bbox head

* add 3dssd detector

* fix bugs for 3dssd config

* modify base conv bbox head

* modify base conv bbox head

* modify base conv bbox head

tests/test_heads.py

@@ -484,7 +484,7 @@ def test_primitive_head():
         num_dims=2,
         num_classes=18,
         primitive_mode='z',
-        vote_moudule_cfg=dict(
+        vote_module_cfg=dict(
             in_channels=256,
             vote_per_seed=1,
             gt_per_seed=1,
@@ -583,7 +583,7 @@ def test_primitive_head():
 
     # 'Primitive_mode' should be one of ['z', 'xy', 'line']
     with pytest.raises(AssertionError):
-        primitive_head_cfg['vote_moudule_cfg']['in_channels'] = 'xyz'
+        primitive_head_cfg['vote_module_cfg']['in_channels'] = 'xyz'
         build_head(primitive_head_cfg)
 
 
@@ -593,18 +593,27 @@ def test_h3d_head():
     _setup_seed(0)
 
     h3d_head_cfg = _get_roi_head_cfg('h3dnet/h3dnet_8x3_scannet-3d-18class.py')
+
+    num_point = 128
+    num_proposal = 64
+    h3d_head_cfg.primitive_list[0].vote_aggregation_cfg.num_point = num_point
+    h3d_head_cfg.primitive_list[1].vote_aggregation_cfg.num_point = num_point
+    h3d_head_cfg.primitive_list[2].vote_aggregation_cfg.num_point = num_point
+    h3d_head_cfg.bbox_head.num_proposal = num_proposal
     self = build_head(h3d_head_cfg).cuda()
 
     # prepare roi outputs
-    fp_xyz = [torch.rand([1, 1024, 3], dtype=torch.float32).cuda()]
-    hd_features = torch.rand([1, 256, 1024], dtype=torch.float32).cuda()
-    fp_indices = [torch.randint(0, 128, [1, 1024]).cuda()]
-    aggregated_points = torch.rand([1, 256, 3], dtype=torch.float32).cuda()
-    aggregated_features = torch.rand([1, 128, 256], dtype=torch.float32).cuda()
+    fp_xyz = [torch.rand([1, num_point, 3], dtype=torch.float32).cuda()]
+    hd_features = torch.rand([1, 256, num_point], dtype=torch.float32).cuda()
+    fp_indices = [torch.randint(0, 128, [1, num_point]).cuda()]
+    aggregated_points = torch.rand([1, num_proposal, 3],
+                                   dtype=torch.float32).cuda()
+    aggregated_features = torch.rand([1, 128, num_proposal],
+                                     dtype=torch.float32).cuda()
     proposal_list = torch.cat([
-        torch.rand([1, 256, 3], dtype=torch.float32).cuda() * 4 - 2,
-        torch.rand([1, 256, 3], dtype=torch.float32).cuda() * 4,
-        torch.zeros([1, 256, 1]).cuda()
+        torch.rand([1, num_proposal, 3], dtype=torch.float32).cuda() * 4 - 2,
+        torch.rand([1, num_proposal, 3], dtype=torch.float32).cuda() * 4,
+        torch.zeros([1, num_proposal, 1]).cuda()
     ],
                               dim=-1)
 
@@ -625,27 +634,32 @@ def test_h3d_head():
     ]
     gt_labels_3d = torch.randint(0, 18, [1, 4]).cuda()
     gt_labels_3d = [gt_labels_3d[0]]
-    pts_semantic_mask = torch.randint(0, 19, [1, 1024]).cuda()
+    pts_semantic_mask = torch.randint(0, 19, [1, num_point]).cuda()
     pts_semantic_mask = [pts_semantic_mask[0]]
-    pts_instance_mask = torch.randint(0, 4, [1, 1024]).cuda()
+    pts_instance_mask = torch.randint(0, 4, [1, num_point]).cuda()
     pts_instance_mask = [pts_instance_mask[0]]
-    points = torch.rand([1, 1024, 3], dtype=torch.float32).cuda()
+    points = torch.rand([1, num_point, 3], dtype=torch.float32).cuda()
 
     # prepare rpn targets
-    vote_targets = torch.rand([1, 1024, 9], dtype=torch.float32).cuda()
-    vote_target_masks = torch.rand([1, 1024], dtype=torch.float32).cuda()
-    size_class_targets = torch.rand([1, 256],
+    vote_targets = torch.rand([1, num_point, 9], dtype=torch.float32).cuda()
+    vote_target_masks = torch.rand([1, num_point], dtype=torch.float32).cuda()
+    size_class_targets = torch.rand([1, num_proposal],
                                     dtype=torch.float32).cuda().long()
-    size_res_targets = torch.rand([1, 256, 3], dtype=torch.float32).cuda()
-    dir_class_targets = torch.rand([1, 256], dtype=torch.float32).cuda().long()
-    dir_res_targets = torch.rand([1, 256], dtype=torch.float32).cuda()
+    size_res_targets = torch.rand([1, num_proposal, 3],
+                                  dtype=torch.float32).cuda()
+    dir_class_targets = torch.rand([1, num_proposal],
+                                   dtype=torch.float32).cuda().long()
+    dir_res_targets = torch.rand([1, num_proposal], dtype=torch.float32).cuda()
     center_targets = torch.rand([1, 4, 3], dtype=torch.float32).cuda()
-    mask_targets = torch.rand([1, 256], dtype=torch.float32).cuda().long()
+    mask_targets = torch.rand([1, num_proposal],
+                              dtype=torch.float32).cuda().long()
     valid_gt_masks = torch.rand([1, 4], dtype=torch.float32).cuda()
-    objectness_targets = torch.rand([1, 256],
+    objectness_targets = torch.rand([1, num_proposal],
                                     dtype=torch.float32).cuda().long()
-    objectness_weights = torch.rand([1, 256], dtype=torch.float32).cuda()
-    box_loss_weights = torch.rand([1, 256], dtype=torch.float32).cuda()
+    objectness_weights = torch.rand([1, num_proposal],
+                                    dtype=torch.float32).cuda()
+    box_loss_weights = torch.rand([1, num_proposal],
+                                  dtype=torch.float32).cuda()
     valid_gt_weights = torch.rand([1, 4], dtype=torch.float32).cuda()
 
     targets = (vote_targets, vote_target_masks, size_class_targets,
@@ -673,3 +687,65 @@ def test_h3d_head():
     assert ret_dict['sem_loss_z'] >= 0
     assert ret_dict['objectness_loss_optimized'] >= 0
     assert ret_dict['primitive_sem_matching_loss'] >= 0
+
+
+def test_ssd3d_head():
+    if not torch.cuda.is_available():
+        pytest.skip('test requires GPU and torch+cuda')
+    _setup_seed(0)
+    ssd3d_head_cfg = _get_vote_head_cfg('3dssd/3dssd_kitti-3d-car.py')
+    ssd3d_head_cfg.vote_module_cfg.num_points = 64
+    self = build_head(ssd3d_head_cfg).cuda()
+    sa_xyz = [torch.rand([2, 128, 3], dtype=torch.float32).cuda()]
+    sa_features = [torch.rand([2, 256, 128], dtype=torch.float32).cuda()]
+    sa_indices = [torch.randint(0, 64, [2, 128]).cuda()]
+
+    input_dict = dict(
+        sa_xyz=sa_xyz, sa_features=sa_features, sa_indices=sa_indices)
+
+    # test forward
+    ret_dict = self(input_dict, 'spec')
+    assert ret_dict['center'].shape == torch.Size([2, 64, 3])
+    assert ret_dict['obj_scores'].shape == torch.Size([2, 1, 64])
+    assert ret_dict['size'].shape == torch.Size([2, 64, 3])
+    assert ret_dict['dir_res'].shape == torch.Size([2, 64, 12])
+
+    # test loss
+    points = [torch.rand([4000, 4], device='cuda') for i in range(2)]
+    gt_bbox1 = LiDARInstance3DBoxes(torch.rand([5, 7], device='cuda'))
+    gt_bbox2 = LiDARInstance3DBoxes(torch.rand([5, 7], device='cuda'))
+    gt_bboxes = [gt_bbox1, gt_bbox2]
+    gt_labels = [
+        torch.zeros([5], dtype=torch.long, device='cuda') for i in range(2)
+    ]
+    img_metas = [dict(box_type_3d=LiDARInstance3DBoxes) for i in range(2)]
+    losses = self.loss(
+        ret_dict, points, gt_bboxes, gt_labels, img_metas=img_metas)
+
+    assert losses['centerness_loss'] >= 0
+    assert losses['center_loss'] >= 0
+    assert losses['dir_class_loss'] >= 0
+    assert losses['dir_res_loss'] >= 0
+    assert losses['size_res_loss'] >= 0
+    assert losses['corner_loss'] >= 0
+    assert losses['vote_loss'] >= 0
+
+    # test multiclass_nms_single
+    sem_scores = ret_dict['obj_scores'].transpose(1, 2)[0]
+    obj_scores = sem_scores.max(-1)[0]
+    bbox = self.bbox_coder.decode(ret_dict)[0]
+    input_meta = img_metas[0]
+    bbox_selected, score_selected, labels = self.multiclass_nms_single(
+        obj_scores, sem_scores, bbox, points[0], input_meta)
+    assert bbox_selected.shape[0] >= 0
+    assert bbox_selected.shape[1] == 7
+    assert score_selected.shape[0] >= 0
+    assert labels.shape[0] >= 0
+
+    # test get_boxes
+    points = torch.stack(points, 0)
+    results = self.get_bboxes(points, ret_dict, img_metas)
+    assert results[0][0].tensor.shape[0] >= 0
+    assert results[0][0].tensor.shape[1] == 7
+    assert results[0][1].shape[0] >= 0
+    assert results[0][2].shape[0] >= 0

tests/test_vote_module.py

@@ -16,6 +16,23 @@ def test_vote_module():
         [2, 8, 64], dtype=torch.float32)  # (b, in_channels, npoints)
 
     # test forward
-    vote_xyz, vote_features = self(seed_xyz, seed_features)
+    vote_xyz, vote_features, vote_offset = self(seed_xyz, seed_features)
     assert vote_xyz.shape == torch.Size([2, 192, 3])
     assert vote_features.shape == torch.Size([2, 8, 192])
+    assert vote_offset.shape == torch.Size([2, 3, 192])
+
+    # test clip offset and without feature residual
+    self = VoteModule(
+        vote_per_seed=1,
+        in_channels=8,
+        num_points=32,
+        with_res_feat=False,
+        vote_xyz_range=(2.0, 2.0, 2.0))
+
+    vote_xyz, vote_features, vote_offset = self(seed_xyz, seed_features)
+    assert vote_xyz.shape == torch.Size([2, 32, 3])
+    assert vote_features.shape == torch.Size([2, 8, 32])
+    assert vote_offset.shape == torch.Size([2, 3, 32])
+    assert torch.allclose(seed_features[..., :32], vote_features)
+    assert vote_offset.max() <= 2.0
+    assert vote_offset.min() >= -2.0