Voxel down sampling operation for point clouds

[Chen-Suyi]

🚀 Feature

pytorch.ops.voxel_down_sample(input, voxel_size) would be a function to downsample input point cloud into output point cloud with a voxel

NOTE: Please look at the existing list of Issues tagged with the label 'enhancement`. Only open a new issue if you do not see your feature request there.

Motivation

We can do voxel down sampling easily by using one of the several libraries, open3d for example. However, most of them are implemented on CPUs and are not compatible with torch.Tensor. It would cost a lot if we do the conversion between torch.Tensor and numpy.ndarray frequently during training. It would be great if we could do this with torch.Tensor on GPUs directly. Some people make their own CUDA versions to do voxel down sampling, but most of the codes are not stable nor compatible with different Pytorch versions. So, It would be helpful if pytorch3d could offer this feature.

Pitch

Voxel down sampling is widely used in various tasks for point clouds.
There are many opensource implementations on the internet. Here for example.

NOTE: we only consider adding new features if they are useful for many users.

[bottler]

The implementation in open3d is https://github.com/isl-org/Open3D/blob/master/cpp/open3d/geometry/PointCloud.cpp#L354 .The idea is to aggregate the points into voxels of a given size, and then return one point per occupied voxel.

It would be quite easy to do a slow implementation of this using PyTorch3D - by adding the pointcloud locations, and normals if present, to the features of the pointcloud; then converting the pointcloud to Volumes, and then extracting a pointcloud from the Volumes and filtering out the unoccupied points. I think this would be quicker than converting to numpy.

The open3d implementation just uses the nearest voxel to each point - i.e. using splatting when converting to Volumes. I am guessing that for training, trilinear would be useful as an option, so that the result has a gradient wrt the input locations? (The downside is more points in the output.)

It would be possible to add a cuda operation of the full thing to pytorch3d.

[Chen-Suyi]

The open3d implementation just uses the nearest voxel to each point - i.e. using splatting when converting to Volumes. I am guessing that for training, trilinear would be useful as an option, so that the result has a gradient wrt the input locations? (The downside is more points in the output.)

To my best knowledge, the gradient of voxel down-sampling is unnecessary most of the time, but it would be great if there was an option to trace the gradient when doing this. I think the gradient would be useful for some tasks in the future.