backward support for spconv VoxelGenerator

pcdet/datasets/processor/data_processor.py

 from functools import partial
 
-import cumm.tensorview as tv
 import numpy as np
 from skimage import transform
-from spconv.utils import Point2VoxelCPU3d as VoxelGenerator
 
 from ...utils import box_utils, common_utils
 
+tv = None
+try:
+    import cumm.tensorview as tv
+except:
+    pass
+
+
+class VoxelGeneratorWrapper():
+    def __init__(self, vsize_xyz, coors_range_xyz, num_point_features, max_num_points_per_voxel, max_num_voxels):
+        try:
+            from spconv.utils import VoxelGeneratorV2 as VoxelGenerator
+            self.spconv_ver = 1
+        except:
+            try:
+                from spconv.utils import VoxelGenerator
+                self.spconv_ver = 1
+            except:
+                from spconv.utils import Point2VoxelCPU3d as VoxelGenerator
+                self.spconv_ver = 2
+
+        if self.spconv_ver == 1:
+            self._voxel_generator = VoxelGenerator(
+                voxel_size=vsize_xyz,
+                point_cloud_range=coors_range_xyz,
+                max_num_points=max_num_points_per_voxel,
+                max_voxels=max_num_voxels
+            )
+        else:
+            self._voxel_generator = VoxelGenerator(
+                vsize_xyz=vsize_xyz,
+                coors_range_xyz=coors_range_xyz,
+                num_point_features=num_point_features,
+                max_num_points_per_voxel=max_num_points_per_voxel,
+                max_num_voxels=max_num_voxels
+            )
+
+    def generate(self, points):
+        if self.spconv_ver == 1:
+            voxel_output = self._voxel_generator.generate(points)
+            if isinstance(voxel_output, dict):
+                voxels, coordinates, num_points = \
+                    voxel_output['voxels'], voxel_output['coordinates'], voxel_output['num_points_per_voxel']
+            else:
+                voxels, coordinates, num_points = voxel_output
+        else:
+            assert tv is not None, f"Unexpected error, library: 'cumm' wasn't imported properly."
+            voxel_output = self._voxel_generator.point_to_voxel(tv.from_numpy(points))
+            tv_voxels, tv_coordinates, tv_num_points = voxel_output
+            # make copy with numpy(), since numpy_view() will disappear as soon as the generator is deleted
+            voxels = tv_voxels.numpy()
+            coordinates = tv_coordinates.numpy()
+            num_points = tv_num_points.numpy()
+        return voxels, coordinates, num_points
+
 
 class DataProcessor(object):
     def __init__(self, processor_configs, point_cloud_range, training, num_point_features):
@@ -16,6 +68,9 @@ class DataProcessor(object):
         self.mode = 'train' if training else 'test'
         self.grid_size = self.voxel_size = None
         self.data_processor_queue = []
+
+        self.voxel_generator = None
+
         for cur_cfg in processor_configs:
             cur_processor = getattr(self, cur_cfg.NAME)(config=cur_cfg)
             self.data_processor_queue.append(cur_processor)
@@ -47,27 +102,27 @@ class DataProcessor(object):
 
         return data_dict
 
-    def transform_points_to_voxels(self, data_dict=None, config=None, voxel_generator=None):
+    def transform_points_to_voxels(self, data_dict=None, config=None):
         if data_dict is None:
-            voxel_generator = VoxelGenerator(
+            grid_size = (self.point_cloud_range[3:6] - self.point_cloud_range[0:3]) / np.array(config.VOXEL_SIZE)
+            self.grid_size = np.round(grid_size).astype(np.int64)
+            self.voxel_size = config.VOXEL_SIZE
+            # just bind the config, we will create the VoxelGeneratorWrapper later,
+            # to avoid pickling issues in multiprocess spawn
+            return partial(self.transform_points_to_voxels, config=config)
+
+        if self.voxel_generator is None:
+            self.voxel_generator = VoxelGeneratorWrapper(
                 vsize_xyz=config.VOXEL_SIZE,
                 coors_range_xyz=self.point_cloud_range,
                 num_point_features=self.num_point_features,
                 max_num_points_per_voxel=config.MAX_POINTS_PER_VOXEL,
-                max_num_voxels=config.MAX_NUMBER_OF_VOXELS[self.mode]
+                max_num_voxels=config.MAX_NUMBER_OF_VOXELS[self.mode],
             )
-            grid_size = (self.point_cloud_range[3:6] - self.point_cloud_range[0:3]) / np.array(config.VOXEL_SIZE)
-            self.grid_size = np.round(grid_size).astype(np.int64)
-            self.voxel_size = config.VOXEL_SIZE
-            return partial(self.transform_points_to_voxels, voxel_generator=voxel_generator)
 
         points = data_dict['points']
-        voxel_output = voxel_generator.point_to_voxel(tv.from_numpy(points))
-        tv_voxels, tv_coordinates, tv_num_points = voxel_output
-        # make copy with numpy(), since numpy_view() will disappear as soon as the generator is deleted
-        voxels = tv_voxels.numpy()
-        coordinates = tv_coordinates.numpy()
-        num_points = tv_num_points.numpy()
+        voxel_output = self.voxel_generator.generate(points)
+        voxels, coordinates, num_points = voxel_output
 
         if not data_dict['use_lead_xyz']:
             voxels = voxels[..., 3:]  # remove xyz in voxels(N, 3)