[Fix] Support inference smoke on Nuscenes and fix bugs (#1029)

* fix collect3d meta keys

* support inference smoke on nuscenes and fix some bugs

* fix some bugs

mmdet3d/core/bbox/coders/smoke_bbox_coder.py

@@ -66,7 +66,7 @@ class SMOKECoder(BaseBBoxCoder):
                 center2d point.
                 shape: (batch, K (max_objs))
             cam2imgs (Tensor): Batch images' camera intrinsic matrix.
-                shape: (batch, 4, 4)
+                shape: kitti (batch, 4, 4)  nuscenes (batch, 3, 3)
             trans_mats (Tensor): transformation matrix from original image
                 to feature map.
                 shape: (batch, 3, 3)
@@ -123,7 +123,7 @@ class SMOKECoder(BaseBBoxCoder):
             depths (Tensor): Object depth z.
                 shape: (batch * K)
             cam2imgs (Tensor): Batch camera intrinsics matrix.
-                shape: (batch, 4, 4)
+                shape: kitti (batch, 4, 4)  nuscenes (batch, 3, 3)
             trans_mats (Tensor): transformation matrix from original image
                 to feature map.
                 shape: (batch, 3, 3)
@@ -144,9 +144,10 @@ class SMOKECoder(BaseBBoxCoder):
         # transform project points back on original image
         centers2d_img = torch.matmul(trans_mats_inv, centers2d_extend)
         centers2d_img = centers2d_img * depths.view(N, -1, 1)
-        centers2d_img_extend = torch.cat(
+        if cam2imgs.shape[1] == 4:
+            centers2d_img = torch.cat(
                 (centers2d_img, centers2d.new_ones(N, 1, 1)), dim=1)
-        locations = torch.matmul(cam2imgs_inv, centers2d_img_extend).squeeze(2)
+        locations = torch.matmul(cam2imgs_inv, centers2d_img).squeeze(2)
 
         return locations[:, :3]
 
@@ -186,16 +187,16 @@ class SMOKECoder(BaseBBoxCoder):
         alphas = torch.atan(ori_vector[:, 0] / (ori_vector[:, 1] + 1e-7))
 
         # get cosine value positive and negative index.
-        cos_pos_inds = (ori_vector[:, 1] >= 0).nonzero()
-        cos_neg_inds = (ori_vector[:, 1] < 0).nonzero()
+        cos_pos_inds = (ori_vector[:, 1] >= 0).nonzero(as_tuple=False)
+        cos_neg_inds = (ori_vector[:, 1] < 0).nonzero(as_tuple=False)
 
         alphas[cos_pos_inds] -= np.pi / 2
         alphas[cos_neg_inds] += np.pi / 2
         # retrieve object rotation y angle.
         yaws = alphas + rays
 
-        larger_inds = (yaws > np.pi).nonzero()
-        small_inds = (yaws < -np.pi).nonzero()
+        larger_inds = (yaws > np.pi).nonzero(as_tuple=False)
+        small_inds = (yaws < -np.pi).nonzero(as_tuple=False)
 
         if len(larger_inds) != 0:
             yaws[larger_inds] -= 2 * np.pi

mmdet3d/datasets/pipelines/formating.py

@@ -138,7 +138,8 @@ class Collect3D(object):
                    'pcd_horizontal_flip', 'pcd_vertical_flip', 'box_mode_3d',
                    'box_type_3d', 'img_norm_cfg', 'pcd_trans', 'sample_idx',
                    'pcd_scale_factor', 'pcd_rotation', 'pcd_rotation_angle',
-                   'pts_filename', 'transformation_3d_flow')):
+                   'pts_filename', 'transformation_3d_flow', 'trans_mat',
+                   'affine_aug')):
         self.keys = keys
         self.meta_keys = meta_keys
 

mmdet3d/models/dense_heads/smoke_mono3d_head.py

@@ -127,7 +127,7 @@ class SMOKEMono3DHead(AnchorFreeMono3DHead):
         """
         assert len(cls_scores) == len(bbox_preds) == 1
         cam2imgs = torch.stack([
-            cls_scores[0].new_tensor(img_meta['cam_intrinsic'])
+            cls_scores[0].new_tensor(img_meta['cam2img'])
             for img_meta in img_metas
         ])
         trans_mats = torch.stack([
@@ -254,7 +254,7 @@ class SMOKEMono3DHead(AnchorFreeMono3DHead):
         batch, channel = pred_reg.shape[0], pred_reg.shape[1]
         w = pred_reg.shape[3]
         cam2imgs = torch.stack([
-            gt_locations.new_tensor(img_meta['cam_intrinsic'])
+            gt_locations.new_tensor(img_meta['cam2img'])
             for img_meta in img_metas
         ])
         trans_mats = torch.stack([

mmdet3d/models/necks/dla_neck.py

@@ -212,7 +212,7 @@ class DLANeck(BaseModule):
         for i in range(self.end_level - self.start_level):
             outs.append(mlvl_features[i].clone())
         self.ida_up(outs, 0, len(outs))
-        return outs[-1]
+        return [outs[-1]]
 
     def init_weights(self):
         for m in self.modules():

tests/test_models/test_heads/test_heads.py

@@ -428,9 +428,8 @@ def test_smoke_mono3d_head():
     attr_labels = None
     img_metas = [
         dict(
-            cam_intrinsic=[[
-                1260.8474446004698, 0.0, 807.968244525554, 40.1111
-            ], [0.0, 1260.8474446004698, 495.3344268742088, 2.34422],
+            cam2img=[[1260.8474446004698, 0.0, 807.968244525554, 40.1111],
+                     [0.0, 1260.8474446004698, 495.3344268742088, 2.34422],
                      [0.0, 0.0, 1.0, 0.00333333], [0.0, 0.0, 0.0, 1.0]],
             scale_factor=np.array([1., 1., 1., 1.], dtype=np.float32),
             pad_shape=[128, 128],

tests/test_models/test_necks/test_necks.py

@@ -131,4 +131,4 @@ def test_dla_neck():
             for i in range(len(in_channels))
         ]
         outputs = neck(feats)
-        assert outputs.shape == (4, 64, 8, 8)
+        assert outputs[0].shape == (4, 64, 8, 8)