v2.1.12: add pc_voxel_id in voxel gen, fix bug

CHANGELOG.md

 # Changelog
 
+## [2.1.12] - 2021-11-23
+### Added 
+- Add a method for voxel generator to get pc_voxel_id, which is usually used in semantic segmentation
+### Fixed
+- Fix a bug in cuda voxel generater when max_voxels is smaller than real number of voxels
+
 ## [2.1.11] - 2021-11-22
 ### Fixed
 - Fixed a bug Volta kernels (TITAN V, Tesla V100), backward weight kernels use f16 as accumulator. we should use f32.

docs/USAGE.md

@@ -82,6 +82,8 @@ class ExampleNet(nn.Module):
 
 Inverse sparse convolution means "inv" of sparse convolution. the output of inverse convolution contains same indices as input of sparse convolution.
 
+**WARNING** ```SparseInverseConv``` isn't equivalent to ```SparseConvTranspose```. SparseConvTranspose is equivalent to ```ConvTranspose``` in pytorch, but SparseInverseConv isn't.
+
 Inverse convolution usually used in semantic segmentation.
 
 ```Python
@@ -112,8 +114,10 @@ voxel generator in spconv generate indices in **ZYX** order, the params format a
 
 generated indices don't include batch axis, you need to add it by yourself.
 
+see examples/voxel_gen.py for examples.
+
 ```Python
-from spconv.pytorch.utils import PointToVoxel
+from spconv.pytorch.utils import PointToVoxel, gather_features_by_pc_voxel_id
 # this generator generate ZYX indices.
 gen = PointToVoxel(
     vsize_xyz=[0.1, 0.1, 0.1], 
@@ -123,5 +127,14 @@ gen = PointToVoxel(
     max_num_points_per_voxel=5)
 pc = np.random.uniform(-10, 10, size=[1000, 3])
 pc_th = torch.from_numpy(pc)
-voxels, coords, num_points_per_voxel = gen(pc_th)
+voxels, coords, num_points_per_voxel = gen(pc_th, empty_mean=True)
 ```
+
+If you want to get label for every point of your pc, you need to use another function to get pc_voxel_id and gather features from sematic segmentation result:
+```Python
+voxels, coords, num_points_per_voxel, pc_voxel_id = gen.generate_voxel_with_id(pc_th, empty_mean=True)
+seg_features = YourSegNet(...)
+# if voxel id is invalid (point out of range, or no space left in a voxel)
+# features will be zero.
+point_features = gather_features_by_pc_voxel_id(seg_features, pc_voxel_id)
+```
\ No newline at end of file

example/voxel_gen.py

@@ -16,9 +16,94 @@ import numpy as np
 
 from cumm import tensorview as tv
 from spconv.utils import Point2VoxelCPU3d
-from spconv.pytorch.utils import PointToVoxel
+from spconv.pytorch.utils import PointToVoxel, gather_features_by_pc_voxel_id
 import torch
 
+def main_pytorch_voxel_gen():
+    np.random.seed(50051)
+    # voxel gen source code: spconv/csrc/sparse/pointops.py
+    gen = PointToVoxel(vsize_xyz=[0.1, 0.1, 0.1],
+                       coors_range_xyz=[-80, -80, -6, 80, 80, 6],
+                       num_point_features=3,
+                       max_num_voxels=5000,
+                       max_num_points_per_voxel=5)
+
+    pc = np.random.uniform(-4, 4, size=[1000, 3])
+    pc_th = torch.from_numpy(pc)
+    voxels_th, indices_th, num_p_in_vx_th = gen(pc_th)
+    voxels_np = voxels_th.numpy()
+    indices_np = indices_th.numpy()
+    num_p_in_vx_np = num_p_in_vx_th.numpy()
+    print(f"------Raw Voxels {voxels_np.shape[0]}-------")
+    print(voxels_np[0])
+    # run voxel gen and FILL MEAN VALUE to voxel remain
+    voxels_th, indices_th, num_p_in_vx_th = gen(pc_th, empty_mean=True)
+    voxels_np = voxels_th.numpy()
+    indices_np = indices_th.numpy()
+    num_p_in_vx_np = num_p_in_vx_th.numpy()
+    print("------Voxels with mean filled-------")
+    print(voxels_np[0])
+    voxels_th, indices_th, num_p_in_vx_th, pc_voxel_id = gen.generate_voxel_with_id(pc_th, empty_mean=True)
+    print("------Voxel ids for every point-------")
+    print(pc_voxel_id[:10])
+
+
+
+def main_pytorch_voxel_gen_cuda():
+    np.random.seed(50051)
+    # voxel gen source code: spconv/csrc/sparse/pointops.py
+    pc = np.random.uniform(-2, 8, size=[1000, 3]).astype(np.float32)
+
+    for device in [torch.device("cuda:0"), torch.device("cpu:0")]:
+        gen = PointToVoxel(vsize_xyz=[0.25, 0.25, 0.25],
+                        coors_range_xyz=[0, 0, 0, 10, 10, 10],
+                        num_point_features=3,
+                        max_num_voxels=5000,
+                        max_num_points_per_voxel=5,
+                        device=device)
+
+        pc_th = torch.from_numpy(pc).to(device)
+        voxels_th, indices_th, num_p_in_vx_th = gen(pc_th)
+        voxels_np = voxels_th.cpu().numpy()
+        indices_np = indices_th.cpu().numpy()
+        num_p_in_vx_np = num_p_in_vx_th.cpu().numpy()
+        print(f"------{device} Raw Voxels {voxels_np.shape[0]}-------")
+        print(voxels_np[0])
+        # run voxel gen and FILL MEAN VALUE to voxel remain
+        voxels_tv, indices_tv, num_p_in_vx_tv = gen(pc_th, empty_mean=True)
+        voxels_np = voxels_tv.cpu().numpy()
+        indices_np = indices_tv.cpu().numpy()
+        num_p_in_vx_np = num_p_in_vx_tv.cpu().numpy()
+        print(f"------{device} Voxels with mean filled-------")
+        print(voxels_np[0])
+        voxels_th, indices_th, num_p_in_vx_th, pc_voxel_id = gen.generate_voxel_with_id(pc_th, empty_mean=True)
+        print(f"------{device} Reconstruct Indices From Voxel ids for every point-------")
+        indices_th_float = indices_th.float()
+        # we gather indices by voxel_id to see correctness of voxel id.
+        indices_th_voxel_id = gather_features_by_pc_voxel_id(indices_th_float, pc_voxel_id)
+        indices_th_voxel_id_np = indices_th_voxel_id[:10].cpu().numpy()
+        print(pc[:10])
+        print(indices_th_voxel_id_np[:, ::-1] / 4)
+
+def main_gather_features_by_pc_voxel_id():
+    np.random.seed(50051)
+    # voxel gen source code: spconv/csrc/sparse/pointops.py
+    device = torch.device("cuda:0")
+    gen = PointToVoxel(vsize_xyz=[0.25, 0.25, 0.25],
+                       coors_range_xyz=[-10, -10, -10, 10, 10, 10],
+                       num_point_features=3,
+                       max_num_voxels=2000,
+                       max_num_points_per_voxel=5,
+                       device=device)
+
+    pc = np.random.uniform(-8, 8, size=[5000, 3]).astype(np.float32)
+    pc_th = torch.from_numpy(pc).to(device)
+
+    voxels_th, indices_th, num_p_in_vx_th, pc_voxel_id = gen.generate_voxel_with_id(pc_th, empty_mean=True)
+    res_features_from_seg = torch.zeros((voxels_th.shape[0], 128), dtype=torch.float32, device=device)
+    
+    pc_features = gather_features_by_pc_voxel_id(res_features_from_seg, pc_voxel_id)
+    print(pc.shape, pc_features.shape)
 
 def main():
     np.random.seed(50051)
@@ -81,58 +166,26 @@ def main_point_with_features():
     print("------Voxels with mean filled-------")
     print(voxels_np[0])
 
-
-def main_pytorch_voxel_gen():
+def main_cuda():
     np.random.seed(50051)
-    # voxel gen source code: spconv/csrc/sparse/pointops.py
-    gen = PointToVoxel(vsize_xyz=[0.1, 0.1, 0.1],
-                       coors_range_xyz=[-80, -80, -2, 80, 80, 6],
-                       num_point_features=3,
-                       max_num_voxels=5000,
-                       max_num_points_per_voxel=5)
+    from spconv.utils import Point2VoxelGPU3d
 
-    pc = np.random.uniform(-10, 10, size=[1000, 3])
-    pc_th = torch.from_numpy(pc)
-    voxels_th, indices_th, num_p_in_vx_th = gen(pc_th)
-    voxels_np = voxels_th.numpy()
-    indices_np = indices_th.numpy()
-    num_p_in_vx_np = num_p_in_vx_th.numpy()
-    print(f"------Raw Voxels {voxels_np.shape[0]}-------")
-    print(voxels_np[0])
-    # run voxel gen and FILL MEAN VALUE to voxel remain
-    voxels_tv, indices_tv, num_p_in_vx_tv = gen(pc_th, empty_mean=True)
-    voxels_np = voxels_tv.numpy()
-    indices_np = indices_tv.numpy()
-    num_p_in_vx_np = num_p_in_vx_tv.numpy()
-    print("------Voxels with mean filled-------")
-    print(voxels_np[0])
-
-
-def main_pytorch_voxel_gen_cuda():
-    np.random.seed(50051)
     # voxel gen source code: spconv/csrc/sparse/pointops.py
-    device = torch.device("cuda:0")
-    gen = PointToVoxel(vsize_xyz=[0.1, 0.1, 0.1],
-                       coors_range_xyz=[-80, -80, -2, 80, 80, 6],
-                       num_point_features=3,
-                       max_num_voxels=5000,
-                       max_num_points_per_voxel=5,
-                       device=device)
+    gen = Point2VoxelGPU3d(vsize_xyz=[0.1, 0.1, 0.1],
+                           coors_range_xyz=[-80, -80, -2, 80, 80, 6],
+                           num_point_features=3,
+                           max_num_voxels=5000,
+                           max_num_points_per_voxel=5)
 
-    pc = np.random.uniform(-10, 10, size=[1000, 3]).astype(np.float32)
-    pc_th = torch.from_numpy(pc).to(device)
-    voxels_th, indices_th, num_p_in_vx_th = gen(pc_th)
-    voxels_np = voxels_th.cpu().numpy()
-    indices_np = indices_th.cpu().numpy()
-    num_p_in_vx_np = num_p_in_vx_th.cpu().numpy()
-    print(f"------Raw Voxels {voxels_np.shape[0]}-------")
-    print(voxels_np[0])
-    # run voxel gen and FILL MEAN VALUE to voxel remain
-    voxels_tv, indices_tv, num_p_in_vx_tv = gen(pc_th, empty_mean=True)
+    pc = np.random.uniform(-10, 10, size=[100000, 3]).astype(np.float32)
+    pc_tv = tv.from_numpy(pc).cuda()
+    # generate voxels, note that voxels_tv reference to a persistent buffer in generator,
+    # so we can't run it in multi-thread.
+    voxels_tv, indices_tv, num_p_in_vx_tv = gen.point_to_voxel_hash(pc_tv)
     voxels_np = voxels_tv.cpu().numpy()
     indices_np = indices_tv.cpu().numpy()
     num_p_in_vx_np = num_p_in_vx_tv.cpu().numpy()
-    print("------Voxels with mean filled-------")
+    print(f"------CUDA Raw Voxels {voxels_np.shape[0]}-------")
     print(voxels_np[0])
 
 
@@ -141,4 +194,6 @@ if __name__ == "__main__":
     main_point_with_features()
     main_pytorch_voxel_gen()
     if torch.cuda.is_available():
+        main_cuda()
         main_pytorch_voxel_gen_cuda()
+        main_gather_features_by_pc_voxel_id()

spconv/core_cc/csrc/sparse/all/__init__.pyi

@@ -291,7 +291,7 @@ class SpconvOps:
         """
         ...
     @staticmethod
-    def point2voxel_cpu(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], empty_mean: bool = False, clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
+    def point2voxel_cpu(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, pc_voxel_id: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], empty_mean: bool = False, clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
         """
         Args:
             points: 
@@ -299,6 +299,7 @@ class SpconvOps:
             indices: 
             num_per_voxel: 
             densehashdata: 
+            pc_voxel_id: 
             vsize: 
             grid_size: 
             grid_stride: 
@@ -308,7 +309,7 @@ class SpconvOps:
         """
         ...
     @staticmethod
-    def point2voxel_cuda(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, hashdata: Tensor, point_indice_data: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], empty_mean: bool = False, clear_voxels: bool = True, stream_int: int = 0) -> Tuple[Tensor, Tensor, Tensor]: 
+    def point2voxel_cuda(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, hashdata: Tensor, point_indice_data: Tensor, pc_voxel_id: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], empty_mean: bool = False, clear_voxels: bool = True, stream_int: int = 0) -> Tuple[Tensor, Tensor, Tensor]: 
         """
         Args:
             points: 
@@ -317,6 +318,7 @@ class SpconvOps:
             num_per_voxel: 
             hashdata: 
             point_indice_data: 
+            pc_voxel_id: 
             vsize: 
             grid_size: 
             grid_stride: 

spconv/core_cc/csrc/sparse/all/ops1d.pyi

@@ -29,7 +29,7 @@ class Point2Voxel:
         """
         ...
     @staticmethod
-    def point_to_voxel_hash_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, hashdata: Tensor, point_indice_data: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True, empty_mean: bool = False, stream_int: int = 0) -> Tuple[Tensor, Tensor, Tensor]: 
+    def point_to_voxel_hash_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, hashdata: Tensor, point_indice_data: Tensor, points_voxel_id: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True, empty_mean: bool = False, stream_int: int = 0) -> Tuple[Tensor, Tensor, Tensor]: 
         """
         Args:
             points: 
@@ -38,6 +38,7 @@ class Point2Voxel:
             num_per_voxel: 
             hashdata: 
             point_indice_data: 
+            points_voxel_id: 
             vsize: 
             grid_size: 
             grid_stride: 

spconv/core_cc/csrc/sparse/all/ops2d.pyi

@@ -29,7 +29,7 @@ class Point2Voxel:
         """
         ...
     @staticmethod
-    def point_to_voxel_hash_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, hashdata: Tensor, point_indice_data: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True, empty_mean: bool = False, stream_int: int = 0) -> Tuple[Tensor, Tensor, Tensor]: 
+    def point_to_voxel_hash_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, hashdata: Tensor, point_indice_data: Tensor, points_voxel_id: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True, empty_mean: bool = False, stream_int: int = 0) -> Tuple[Tensor, Tensor, Tensor]: 
         """
         Args:
             points: 
@@ -38,6 +38,7 @@ class Point2Voxel:
             num_per_voxel: 
             hashdata: 
             point_indice_data: 
+            points_voxel_id: 
             vsize: 
             grid_size: 
             grid_stride: 

spconv/core_cc/csrc/sparse/all/ops3d.pyi

@@ -29,7 +29,7 @@ class Point2Voxel:
         """
         ...
     @staticmethod
-    def point_to_voxel_hash_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, hashdata: Tensor, point_indice_data: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True, empty_mean: bool = False, stream_int: int = 0) -> Tuple[Tensor, Tensor, Tensor]: 
+    def point_to_voxel_hash_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, hashdata: Tensor, point_indice_data: Tensor, points_voxel_id: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True, empty_mean: bool = False, stream_int: int = 0) -> Tuple[Tensor, Tensor, Tensor]: 
         """
         Args:
             points: 
@@ -38,6 +38,7 @@ class Point2Voxel:
             num_per_voxel: 
             hashdata: 
             point_indice_data: 
+            points_voxel_id: 
             vsize: 
             grid_size: 
             grid_stride: 

spconv/core_cc/csrc/sparse/all/ops4d.pyi

@@ -29,7 +29,7 @@ class Point2Voxel:
         """
         ...
     @staticmethod
-    def point_to_voxel_hash_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, hashdata: Tensor, point_indice_data: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True, empty_mean: bool = False, stream_int: int = 0) -> Tuple[Tensor, Tensor, Tensor]: 
+    def point_to_voxel_hash_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, hashdata: Tensor, point_indice_data: Tensor, points_voxel_id: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True, empty_mean: bool = False, stream_int: int = 0) -> Tuple[Tensor, Tensor, Tensor]: 
         """
         Args:
             points: 
@@ -38,6 +38,7 @@ class Point2Voxel:
             num_per_voxel: 
             hashdata: 
             point_indice_data: 
+            points_voxel_id: 
             vsize: 
             grid_size: 
             grid_stride: 

spconv/core_cc/csrc/sparse/all/ops_cpu1d.pyi

@@ -27,7 +27,7 @@ class Point2VoxelCPU:
         """
         ...
     @staticmethod
-    def point_to_voxel_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
+    def point_to_voxel_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, points_voxel_id: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
         """
         Args:
             points: 
@@ -35,6 +35,7 @@ class Point2VoxelCPU:
             indices: 
             num_per_voxel: 
             densehashdata: 
+            points_voxel_id: 
             vsize: 
             grid_size: 
             grid_stride: 
@@ -43,7 +44,7 @@ class Point2VoxelCPU:
         """
         ...
     @staticmethod
-    def point_to_voxel_empty_mean_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
+    def point_to_voxel_empty_mean_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, points_voxel_id: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
         """
         Args:
             points: 
@@ -51,6 +52,7 @@ class Point2VoxelCPU:
             indices: 
             num_per_voxel: 
             densehashdata: 
+            points_voxel_id: 
             vsize: 
             grid_size: 
             grid_stride: 

spconv/core_cc/csrc/sparse/all/ops_cpu2d.pyi

@@ -27,7 +27,7 @@ class Point2VoxelCPU:
         """
         ...
     @staticmethod
-    def point_to_voxel_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
+    def point_to_voxel_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, points_voxel_id: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
         """
         Args:
             points: 
@@ -35,6 +35,7 @@ class Point2VoxelCPU:
             indices: 
             num_per_voxel: 
             densehashdata: 
+            points_voxel_id: 
             vsize: 
             grid_size: 
             grid_stride: 
@@ -43,7 +44,7 @@ class Point2VoxelCPU:
         """
         ...
     @staticmethod
-    def point_to_voxel_empty_mean_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
+    def point_to_voxel_empty_mean_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, points_voxel_id: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
         """
         Args:
             points: 
@@ -51,6 +52,7 @@ class Point2VoxelCPU:
             indices: 
             num_per_voxel: 
             densehashdata: 
+            points_voxel_id: 
             vsize: 
             grid_size: 
             grid_stride: 

spconv/core_cc/csrc/sparse/all/ops_cpu3d.pyi

@@ -27,7 +27,7 @@ class Point2VoxelCPU:
         """
         ...
     @staticmethod
-    def point_to_voxel_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
+    def point_to_voxel_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, points_voxel_id: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
         """
         Args:
             points: 
@@ -35,6 +35,7 @@ class Point2VoxelCPU:
             indices: 
             num_per_voxel: 
             densehashdata: 
+            points_voxel_id: 
             vsize: 
             grid_size: 
             grid_stride: 
@@ -43,7 +44,7 @@ class Point2VoxelCPU:
         """
         ...
     @staticmethod
-    def point_to_voxel_empty_mean_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
+    def point_to_voxel_empty_mean_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, points_voxel_id: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
         """
         Args:
             points: 
@@ -51,6 +52,7 @@ class Point2VoxelCPU:
             indices: 
             num_per_voxel: 
             densehashdata: 
+            points_voxel_id: 
             vsize: 
             grid_size: 
             grid_stride: 

spconv/core_cc/csrc/sparse/all/ops_cpu4d.pyi

@@ -27,7 +27,7 @@ class Point2VoxelCPU:
         """
         ...
     @staticmethod
-    def point_to_voxel_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
+    def point_to_voxel_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, points_voxel_id: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
         """
         Args:
             points: 
@@ -35,6 +35,7 @@ class Point2VoxelCPU:
             indices: 
             num_per_voxel: 
             densehashdata: 
+            points_voxel_id: 
             vsize: 
             grid_size: 
             grid_stride: 
@@ -43,7 +44,7 @@ class Point2VoxelCPU:
         """
         ...
     @staticmethod
-    def point_to_voxel_empty_mean_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
+    def point_to_voxel_empty_mean_static(points: Tensor, voxels: Tensor, indices: Tensor, num_per_voxel: Tensor, densehashdata: Tensor, points_voxel_id: Tensor, vsize: List[float], grid_size: List[int], grid_stride: List[int], coors_range: List[float], clear_voxels: bool = True) -> Tuple[Tensor, Tensor, Tensor]: 
         """
         Args:
             points: 
@@ -51,6 +52,7 @@ class Point2VoxelCPU:
             indices: 
             num_per_voxel: 
             densehashdata: 
+            points_voxel_id: 
             vsize: 
             grid_size: 
             grid_stride: 

spconv/csrc/sparse/all.py

@@ -860,7 +860,7 @@ class SpconvOps(pccm.Class):
     def point2voxel_cpu(self):
         code = pccm.FunctionCode()
         code.arg("points", "tv::Tensor")
-        code.arg("voxels, indices, num_per_voxel, densehashdata", "tv::Tensor")
+        code.arg("voxels, indices, num_per_voxel, densehashdata, pc_voxel_id", "tv::Tensor")
         code.arg("vsize", f"std::vector<float>")
         code.arg("grid_size, grid_stride", f"std::vector<int>")
         code.arg("coors_range", f"std::vector<float>")
@@ -890,11 +890,11 @@ class SpconvOps(pccm.Class):
                 }}
                 if (empty_mean){{
                     return Point2Voxel{ndim}DCPU::point_to_voxel_empty_mean_static(points, voxels, indices, 
-                        num_per_voxel, densehashdata, 
+                        num_per_voxel, densehashdata, pc_voxel_id,
                         vsize_, grid_size_, grid_stride_, coors_range_, clear_voxels);
                 }} else{{
                     return Point2Voxel{ndim}DCPU::point_to_voxel_static(points, voxels, indices, 
-                        num_per_voxel, densehashdata, 
+                        num_per_voxel, densehashdata, pc_voxel_id,
                         vsize_, grid_size_, grid_stride_, coors_range_, clear_voxels);
                 }}
             }}
@@ -907,7 +907,7 @@ class SpconvOps(pccm.Class):
     def point2voxel_cuda(self):
         code = pccm.FunctionCode()
         code.arg("points", "tv::Tensor")
-        code.arg("voxels, indices, num_per_voxel, hashdata, point_indice_data",
+        code.arg("voxels, indices, num_per_voxel, hashdata, point_indice_data, pc_voxel_id",
                  "tv::Tensor")
         code.arg("vsize", f"std::vector<float>")
         code.arg("grid_size, grid_stride", f"std::vector<int>")
@@ -940,7 +940,7 @@ class SpconvOps(pccm.Class):
                     coors_range_[i + {ndim}] = coors_range[i + {ndim}];
                 }}
                 return Point2Voxel{ndim}D::point_to_voxel_hash_static(points, voxels, indices, 
-                    num_per_voxel, hashdata, point_indice_data,
+                    num_per_voxel, hashdata, point_indice_data, pc_voxel_id,
                     vsize_, grid_size_, grid_stride_, coors_range_, clear_voxels, 
                     empty_mean, stream_int);
             }}

spconv/csrc/sparse/pointops.py

@@ -208,6 +208,7 @@ class Point2VoxelKernel(pccm.ParameterizedClass, pccm.pybind.PybindClassMixin):
         code.arg("points_indice_data", f"const int64_t*")
         code.arg("voxels", f"{self.dtype} *")
         code.arg("num_per_voxel", f"int *")
+        code.arg("points_voxel_id", f"int64_t*")
 
         code.arg("point_stride", f"int")
         code.arg("max_points_per_voxel", f"int")
@@ -219,14 +220,17 @@ class Point2VoxelKernel(pccm.ParameterizedClass, pccm.pybind.PybindClassMixin):
         code.arg("grid_stride", f"tv::array<int, {self.ndim}>")
 
         code.arg("num_points", f"int")
+        # TODO add backward?
         code.raw(f"""
         int voxel_stride0 = point_stride * max_points_per_voxel;
         for (int i : tv::KernelLoopX<int>(num_points)){{
             int64_t prod = points_indice_data[i];
+            int voxel_id = -1;
             if (prod != -1){{
                 auto voxel_index_pair = table.lookup(prod);
                 if (!voxel_index_pair.empty() &&
                     voxel_index_pair.second < max_voxels) {{
+                    voxel_id = voxel_index_pair.second;
                     int old = atomicAdd(num_per_voxel + voxel_index_pair.second, 1);
                     if (old < max_points_per_voxel) {{
                         for (int j = 0; j < point_stride; ++j) {{
@@ -235,6 +239,7 @@ class Point2VoxelKernel(pccm.ParameterizedClass, pccm.pybind.PybindClassMixin):
                     }}
                 }}
             }}
+            points_voxel_id[i] = voxel_id;
         }}
         """)
         return code
@@ -385,6 +390,7 @@ class Point2Voxel(pccm.ParameterizedClass, pccm.pybind.PybindClassMixin):
         code.arg("stream_int", f"std::uintptr_t", "0")
 
         code.raw(f"""
+        tv::Tensor points_voxel_id = tv::empty({{points.dim(0)}}, tv::int64, 0);
         int64_t expected_hash_data_num = points.dim(0) * 2;
         if (hashdata.dim(0) < expected_hash_data_num){{
             hashdata = tv::zeros({{expected_hash_data_num}}, tv::custom128, 0);
@@ -393,74 +399,18 @@ class Point2Voxel(pccm.ParameterizedClass, pccm.pybind.PybindClassMixin):
             point_indice_data = tv::zeros({{points.dim(0)}}, tv::int64, 0);
         }}
         return point_to_voxel_hash_static(points, voxels, indices, num_per_voxel, 
-            hashdata, point_indice_data, Point2VoxelCommon::tvarray2array(vsize), 
+            hashdata, point_indice_data, points_voxel_id, Point2VoxelCommon::tvarray2array(vsize), 
             Point2VoxelCommon::tvarray2array(grid_size), Point2VoxelCommon::tvarray2array(grid_stride), 
             Point2VoxelCommon::tvarray2array(coors_range), clear_voxels, empty_mean, stream_int);
         """)
         return code.ret("std::tuple<tv::Tensor, tv::Tensor, tv::Tensor>")
-        code.raw(f"""
-
-        TV_ASSERT_INVALID_ARG(points.ndim() == 2 && points.dim(1) >= {self.ndim}, "error");
-        using V = int64_t;
-        using KeyType = int64_t;
-        constexpr KeyType kEmptyKey = std::numeric_limits<KeyType>::max();
-        if (clear_voxels){{
-            voxels.zero_();
-        }}
-        using table_t =
-            tv::hash::LinearHashTable<KeyType, V, tv::hash::Murmur3Hash<KeyType>,
-                                        kEmptyKey, false>;
-        using pair_t = typename table_t::value_type;
-        // int64_t expected_hash_data_num = int64_t(tv::hash::align_to_power2(points.dim(0) * 2));
-        int64_t expected_hash_data_num = points.dim(0) * 2;
-
-        if (hashdata.dim(0) < expected_hash_data_num){{
-            hashdata = tv::zeros({{expected_hash_data_num}}, tv::custom128, 0);
-        }}
-        if (point_indice_data.dim(0) < points.dim(0)){{
-            point_indice_data = tv::zeros({{points.dim(0)}}, tv::int64, 0);
-        }}
-        // auto timer = tv::CudaContextTimer<>();
-        num_per_voxel.zero_();
-        table_t hash = table_t(hashdata.data_ptr<pair_t>(), expected_hash_data_num);
-        hash.clear();
-        // tv::ssprint("clear time", timer.report());
-        auto launcher = tv::cuda::Launch(points.dim(0));
-        launcher(kernel::build_hash_table<table_t>, hash, points.data_ptr<const {self.dtype}>(),
-                point_indice_data.data_ptr<int64_t>(),
-                points.dim(1), vsize, coors_range, grid_size, grid_stride, points.dim(0));
-        // tv::ssprint("build_hash_table", timer.report());
-
-        auto table_launcher = tv::cuda::Launch(hash.size());
-        tv::Tensor count = tv::zeros({{1}}, tv::int32, 0);
-        Layout layout = Layout::from_shape(grid_size);
-        table_launcher(kernel::assign_table<table_t>, hash, indices.data_ptr<int>(),
-                        count.data_ptr<int>(),
-                        layout, voxels.dim(0));
-        auto count_cpu = count.cpu();
-        int count_val = count_cpu.item<int32_t>();
-        // tv::ssprint("assign_table", timer.report());
-
-        launcher(kernel::generate_voxel<table_t>, hash, points.data_ptr<const {self.dtype}>(),
-                point_indice_data.data_ptr<const int64_t>(), voxels.data_ptr<{self.dtype}>(),
-                num_per_voxel.data_ptr<int>(), points.dim(1), voxels.dim(1), 
-                voxels.dim(0), vsize, coors_range,
-                grid_size, grid_stride, points.dim(0));
-        // tv::ssprint("generate_voxel", timer.report());
-
-        return std::make_tuple(voxels.slice_first_axis(0, count_val), 
-            indices.slice_first_axis(0, count_val), 
-            num_per_voxel.slice_first_axis(0, count_val));
-
-        """)
-        return code.ret("std::tuple<tv::Tensor, tv::Tensor, tv::Tensor>")
 
     @pccm.pybind.mark
     @pccm.cuda.static_function
     def point_to_voxel_hash_static(self):
         code = pccm.FunctionCode()
         code.arg("points", "tv::Tensor")
-        code.arg("voxels, indices, num_per_voxel, hashdata, point_indice_data",
+        code.arg("voxels, indices, num_per_voxel, hashdata, point_indice_data, points_voxel_id",
                  "tv::Tensor")
         code.arg("vsize", f"std::array<float, {self.ndim}>")
         code.arg("grid_size, grid_stride", f"std::array<int, {self.ndim}>")
@@ -495,7 +445,6 @@ class Point2Voxel(pccm.ParameterizedClass, pccm.pybind.PybindClassMixin):
         int64_t expected_hash_data_num = points.dim(0) * 2;
         TV_ASSERT_RT_ERR(hashdata.dim(0) >= expected_hash_data_num, "hash table too small")
         TV_ASSERT_RT_ERR(point_indice_data.dim(0) >= points.dim(0), "point_indice_data too small")
-        // auto timer = tv::CudaContextTimer<>();
         num_per_voxel.zero_(ctx);
         table_t hash = table_t(hashdata.data_ptr<pair_t>(), expected_hash_data_num);
         hash.clear(custream);
@@ -512,14 +461,12 @@ class Point2Voxel(pccm.ParameterizedClass, pccm.pybind.PybindClassMixin):
                         layout, voxels.dim(0));
         auto count_cpu = count.cpu();
         int count_val = count_cpu.item<int32_t>();
-        // tv::ssprint("assign_table", timer.report());
-
+        count_val = count_val > voxels.dim(0) ? voxels.dim(0) : count_val;
         launcher(kernel::generate_voxel<table_t>, hash, points.data_ptr<const {self.dtype}>(),
                 point_indice_data.data_ptr<const int64_t>(), voxels.data_ptr<{self.dtype}>(),
-                num_per_voxel.data_ptr<int>(), points.dim(1), voxels.dim(1), 
+                num_per_voxel.data_ptr<int>(), points_voxel_id.data_ptr<int64_t>(), points.dim(1), voxels.dim(1), 
                 voxels.dim(0), vsize_tv, coors_range_tv,
                 grid_size_tv, grid_stride_tv, points.dim(0));
-        // tv::ssprint("generate_voxel", timer.report());
         auto voxel_launcher = tv::cuda::Launch(count_val, custream);
         if (empty_mean){{
             launcher(kernel::voxel_empty_fill_mean, voxels.data_ptr<{self.dtype}>(),
@@ -636,7 +583,7 @@ class Point2VoxelCPU(pccm.ParameterizedClass, pccm.pybind.PybindClassMixin):
     def point_to_voxel_static_template(self, mean: bool = False):
         code = pccm.FunctionCode()
         code.arg("points", "tv::Tensor")
-        code.arg("voxels, indices, num_per_voxel, densehashdata", "tv::Tensor")
+        code.arg("voxels, indices, num_per_voxel, densehashdata, points_voxel_id", "tv::Tensor")
         code.arg("vsize", f"std::array<float, {self.ndim}>")
         code.arg("grid_size, grid_stride", f"std::array<int, {self.ndim}>")
         code.arg("coors_range", f"std::array<float, {self.ndim * 2}>")
@@ -653,6 +600,7 @@ class Point2VoxelCPU(pccm.ParameterizedClass, pccm.pybind.PybindClassMixin):
         if (clear_voxels){{
             voxels.zero_();
         }}
+        auto points_voxel_id_ptr = points_voxel_id.data_ptr<int64_t>();
         int res_voxel_num = 0;
         int num_features = points.dim(1);
         auto N = points.dim(0);
@@ -680,20 +628,25 @@ class Point2VoxelCPU(pccm.ParameterizedClass, pccm.pybind.PybindClassMixin):
                     }}
                     coor[j] = c;
                 }}
-                if (failed)
+                if (failed){{
+                    points_voxel_id_ptr[i] = -1;
                     continue;
+                }}
                 voxelidx = coor_to_voxelidx_rw({codeops.unpack("coor", range(self.ndim))});
-
+                
                 if (voxelidx == -1) {{
                     voxelidx = voxel_num;
-                    if (voxel_num >= max_num_voxels)
+                    if (voxel_num >= max_num_voxels){{
+                        points_voxel_id_ptr[i] = -1;
                         continue;
+                    }}
                     voxel_num += 1;
                     coor_to_voxelidx_rw({codeops.unpack("coor", range(self.ndim))}) = voxelidx;
                     for (int k = 0; k < {self.ndim}; ++k) {{
                         coors_rw(voxelidx, k) = coor[k];
                     }}
                 }}
+                points_voxel_id_ptr[i] = voxelidx;
                 num = num_points_per_voxel_rw(voxelidx);
                 if (num < max_num_points_per_voxel) {{
                     // voxel_point_mask_rw(voxelidx, num) = {self.dtype}(1);
@@ -781,8 +734,10 @@ class Point2VoxelCPU(pccm.ParameterizedClass, pccm.pybind.PybindClassMixin):
         code.arg("points", "tv::Tensor")
         code.arg("clear_voxels", "bool", "true")
         code.raw(f"""
+        tv::Tensor points_voxel_id = tv::empty({{points.dim(0)}}, tv::int64, -1);
+
         return point_to_voxel_static(points, voxels, indices, num_per_voxel, densehashdata, 
-            tvarray2array(vsize), 
+            points_voxel_id, tvarray2array(vsize), 
             tvarray2array(grid_size), tvarray2array(grid_stride), 
             tvarray2array(coors_range), clear_voxels);
         """)
@@ -795,8 +750,10 @@ class Point2VoxelCPU(pccm.ParameterizedClass, pccm.pybind.PybindClassMixin):
         code.arg("points", "tv::Tensor")
         code.arg("clear_voxels", "bool", "true")
         code.raw(f"""
+        tv::Tensor points_voxel_id = tv::empty({{points.dim(0)}}, tv::int64, -1);
+
         return point_to_voxel_empty_mean_static(points, voxels, indices, num_per_voxel, 
-            densehashdata, tvarray2array(vsize), 
+            densehashdata, points_voxel_id, tvarray2array(vsize), 
             tvarray2array(grid_size), tvarray2array(grid_stride), 
             tvarray2array(coors_range), clear_voxels);
         """)

spconv/pytorch/utils.py

@@ -71,45 +71,83 @@ class PointToVoxel(object):
                  pc: torch.Tensor,
                  clear_voxels: bool = True,
                  empty_mean: bool = False):
+        """generate voxels/indices/num_point_per_voxel/pc_voxel_ids from 
+        point cloud.
+        This function don't return pc_voxel_id for backward compatility.
+        pc_voxel_id will be added in spconv 2.2.
+        Args:
+            pc: [N, 3+] point cloud.
+            clear_voxels: if True, call zero on voxels
+            empty_mean: if True, full empty location of voxels with mean.
+        Returns:
+            voxels: voxels
+            indices: quantized coords
+            num_per_voxel: number of points in a voxel
+        """
+
+        res = self.generate_voxel_with_id(pc, clear_voxels, empty_mean)
+        return res[0], res[1], res[2]
+
+    def generate_voxel_with_id(self,
+                 pc: torch.Tensor,
+                 clear_voxels: bool = True,
+                 empty_mean: bool = False):
+        """generate voxels/indices/num_point_per_voxel/pc_voxel_ids from 
+        point cloud.
+        Args:
+            pc: [N, 3+] point cloud.
+            clear_voxels: if True, call zero on voxels
+            empty_mean: if True, full empty location of voxels with mean.
+        Returns:
+            voxels: voxels
+            indices: quantized coords
+            num_per_voxel: number of points in a voxel
+            pc_voxel_id: voxel id for every point. if not exists, -1.
+        """
         assert pc.device.type == self.device.type, "your pc device is wrong"
         expected_hash_data_num = pc.shape[0] * 2
         with torch.no_grad():
+            pc_voxel_id = torch.empty([pc.shape[0]],
+                                    dtype=torch.int64,
+                                    device=self.device)
+            pc_voxel_id_tv = torch_tensor_to_tv(pc_voxel_id)
+
             if self.device.type != "cpu":
-                if self.hashdata.shape[0] < expected_hash_data_num:
-                    self.hashdata = torch.empty([expected_hash_data_num, 2],
-                                                dtype=torch.int64,
-                                                device=self.device)
-
-                if self.point_indice_data.shape[0] < pc.shape[0]:
-                    self.point_indice_data = torch.empty([pc.shape[0]],
-                                                         dtype=torch.int64,
-                                                         device=self.device)
+                hashdata = torch.empty([expected_hash_data_num, 2],
+                                            dtype=torch.int64,
+                                            device=pc.device)
+
+                point_indice_data = torch.empty([pc.shape[0]],
+                                            dtype=torch.int64,
+                                            device=pc.device)
+
                 pc_tv = torch_tensor_to_tv(pc)
                 stream = get_current_stream()
                 voxels_tv = torch_tensor_to_tv(self.voxels)
                 indices_tv = torch_tensor_to_tv(self.indices)
                 num_per_voxel_tv = torch_tensor_to_tv(self.num_per_voxel)
                 hashdata_tv = torch_tensor_to_tv(
-                    self.hashdata,
+                    hashdata,
                     dtype=tv.custom128,
-                    shape=[self.hashdata.shape[0]])
-                point_indice_data_tv = torch_tensor_to_tv(
-                    self.point_indice_data)
-
-                res = SpconvOps.point2voxel_cuda(
-                    pc_tv, voxels_tv, indices_tv, num_per_voxel_tv,
-                    hashdata_tv, point_indice_data_tv, self.vsize,
-                    self.grid_size, self.grid_stride, self.coors_range,
-                    empty_mean, clear_voxels, stream)
+                    shape=[hashdata.shape[0]])
+                point_indice_data_tv = torch_tensor_to_tv(point_indice_data)
+                with torch.cuda.device(pc.device):
+                    res = SpconvOps.point2voxel_cuda(
+                        pc_tv, voxels_tv, indices_tv, num_per_voxel_tv,
+                        hashdata_tv, point_indice_data_tv, pc_voxel_id_tv, self.vsize,
+                        self.grid_size, self.grid_stride, self.coors_range,
+                        empty_mean, clear_voxels, stream)
                 num_voxels = res[0].shape[0]
             else:
                 pc_tv = torch_tensor_to_tv(pc)
+                stream = get_current_stream()
                 voxels_tv = torch_tensor_to_tv(self.voxels)
                 indices_tv = torch_tensor_to_tv(self.indices)
                 num_per_voxel_tv = torch_tensor_to_tv(self.num_per_voxel)
                 hashdata_tv = torch_tensor_to_tv(self.hashdata, dtype=tv.int32)
                 res = SpconvOps.point2voxel_cpu(pc_tv, voxels_tv, indices_tv,
                                                 num_per_voxel_tv, hashdata_tv,
+                                                pc_voxel_id_tv,
                                                 self.vsize, self.grid_size,
                                                 self.grid_stride,
                                                 self.coors_range, empty_mean,
@@ -117,4 +155,18 @@ class PointToVoxel(object):
                 num_voxels = res[0].shape[0]
 
             return (self.voxels[:num_voxels], self.indices[:num_voxels],
-                    self.num_per_voxel[:num_voxels])
+                    self.num_per_voxel[:num_voxels], pc_voxel_id)
+
+
+def gather_features_by_pc_voxel_id(seg_res_features: torch.Tensor, pc_voxel_id: torch.Tensor):
+    """This function is used to gather segmentation result to match origin pc.
+    """
+    if seg_res_features.device != pc_voxel_id.device:
+        pc_voxel_id = pc_voxel_id.to(seg_res_features.device)
+    res = torch.zeros((pc_voxel_id.shape[0], seg_res_features.shape[1]), dtype=seg_res_features.dtype, device=seg_res_features.device)
+    pc_voxel_id_valid = pc_voxel_id != -1
+    pc_voxel_id_valid_ids = torch.nonzero(pc_voxel_id_valid).view(-1)
+    seg_res_features_valid = seg_res_features[pc_voxel_id[pc_voxel_id_valid_ids]]
+    res[pc_voxel_id_valid_ids] = seg_res_features_valid
+    return res 
+

version.txt

-2.1.11
+2.1.12