pass voxel size and pcrange to PillarFeatureNet

Liveshkumar · Mar 18, 2019 · 64ac6a8 · 64ac6a8
1 parent e8ada4e
commit 64ac6a8
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/second/pytorch/builder/second_builder.py b/second/pytorch/builder/second_builder.py
@@ -90,5 +90,7 @@ def build(model_cfg: second_pb2.VoxelNet, voxel_generator,
         loc_loss_ftor=loc_loss_ftor,
         cls_loss_ftor=cls_loss_ftor,
         target_assigner=target_assigner,
+        voxel_size=voxel_generator.voxel_size,
+        pc_range=voxel_generator.point_cloud_range
     )
     return net
diff --git a/second/pytorch/models/voxelnet.py b/second/pytorch/models/voxelnet.py
@@ -528,6 +528,8 @@ def __init__(self,
                  encode_rad_error_by_sin=False,
                  loc_loss_ftor=None,
                  cls_loss_ftor=None,
+                 voxel_size=(0.2, 0.2, 4),
+                 pc_range=(0, -40, -3, 70.4, 40, 1),
                  name='voxelnet'):
         super().__init__()
         self.name = name
@@ -568,11 +570,21 @@ def __init__(self,
             "PillarFeatureNet": PillarFeatureNet
         }
         vfe_class = vfe_class_dict[vfe_class_name]
-        self.voxel_feature_extractor = vfe_class(
-            num_input_features,
-            use_norm,
-            num_filters=vfe_num_filters,
-            with_distance=with_distance)
+        if vfe_class_name == "PillarFeatureNet":
+            self.voxel_feature_extractor = vfe_class(
+                num_input_features,
+                use_norm,
+                num_filters=vfe_num_filters,
+                with_distance=with_distance,
+                voxel_size=voxel_size,
+                pc_range=pc_range
+            )
+        else:
+            self.voxel_feature_extractor = vfe_class(
+                num_input_features,
+                use_norm,
+                num_filters=vfe_num_filters,
+                with_distance=with_distance)
 
         print("middle_class_name", middle_class_name)
         if middle_class_name == "PointPillarsScatter":