fix pillarencode (open-mmlab#331)

humu789 · May 5, 2022 · 86ab063 · 86ab063
1 parent 16ee9c7
commit 86ab063
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 12 deletions.
diff --git a/mmdeploy/codebase/mmdet3d/deploy/voxel_detection_model.py b/mmdeploy/codebase/mmdet3d/deploy/voxel_detection_model.py
@@ -174,7 +174,7 @@ def voxelize(model_cfg: Union[str, mmcv.Config], points: torch.Tensor):
     @staticmethod
     def post_process(model_cfg: Union[str, mmcv.Config],
                      deploy_cfg: Union[str, mmcv.Config],
-                     outs: torch.Tensor,
+                     outs: Dict,
                      img_metas: Dict,
                      device: str,
                      rescale=False):
@@ -184,7 +184,7 @@ def post_process(model_cfg: Union[str, mmcv.Config],
             model_cfg (str | mmcv.Config): The model config.
             deploy_cfg (str|mmcv.Config): Deployment config file or loaded
             Config object.
-            outs (torch.Tensor): Output of model's head.
+            outs (Dict): Output of model's head.
             img_metas(Dict): Meta info for pcd.
             device (str): A string specifying device type.
             rescale (list[torch.Tensor]): whether th rescale bbox.

diff --git a/mmdeploy/codebase/mmdet3d/models/pillar_encode.py b/mmdeploy/codebase/mmdet3d/models/pillar_encode.py
@@ -30,19 +30,18 @@ def pillar_encoder__forward(ctx, self, features, num_points, coors):
 
     # Find distance of x, y, and z from pillar center
     device = features.device
+
     if self._with_voxel_center:
         if not self.legacy:
-            f_center = features[..., :3] - (
-                coors * torch.tensor([1, self.vz, self.vy, self.vx]).to(device)
-                +
-                torch.tensor([1, self.z_offset, self.y_offset, self.x_offset
-                              ]).to(device)).unsqueeze(1).flip(2)[..., :3]
+            f_center = features[..., :3] - (coors[..., 1:] * torch.tensor(
+                [self.vz, self.vy, self.vx]).to(device) + torch.tensor([
+                    self.z_offset, self.y_offset, self.x_offset
+                ]).to(device)).unsqueeze(1).flip(2)
         else:
-            f_center = features[..., :3] - (
-                coors * torch.tensor([1, self.vz, self.vy, self.vx]).to(device)
-                +
-                torch.tensor([1, self.z_offset, self.y_offset, self.x_offset
-                              ]).to(device)).unsqueeze(1).flip(2)[..., :3]
+            f_center = features[..., :3] - (coors[..., 1:] * torch.tensor(
+                [self.vz, self.vy, self.vx]).to(device) + torch.tensor([
+                    self.z_offset, self.y_offset, self.x_offset
+                ]).to(device)).unsqueeze(1).flip(2)
             features_ls[0] = torch.cat((f_center, features[..., 3:]), dim=-1)
         features_ls.append(f_center)