Create Python API and test cases for Cubic Feature Sampling.

Author: hzxie

test/test_cubic_feature_sampling.py

@@ -0,0 +1,60 @@
+import numpy as np
+import os
+import sys
+import torch
+import unittest
+
+from torch.autograd import gradcheck
+
+from . import run_if_cuda
+
+
+ROOT = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")
+sys.path.insert(0, ROOT)
+
+from torch_points_kernels import CubicFeatureSamplingFunction, cubic_feature_sampling
+
+
+class TestCubicFeatureSampling(unittest.TestCase):
+    @run_if_cuda
+    def test_neighborhood_size_1(self):
+        ptcloud = torch.rand(2, 64, 3) * 2 - 1
+        cubic_features = torch.rand(2, 4, 8, 8, 8)
+        ptcloud.requires_grad = True
+        cubic_features.requires_grad = True
+        self.assertTrue(
+            gradcheck(
+                CubicFeatureSamplingFunction.apply,
+                [ptcloud.double().cuda(), cubic_features.double().cuda()],
+            )
+        )
+
+    @run_if_cuda
+    def test_neighborhood_size_2(self):
+        ptcloud = torch.rand(2, 32, 3) * 2 - 1
+        cubic_features = torch.rand(2, 2, 8, 8, 8)
+        ptcloud.requires_grad = True
+        cubic_features.requires_grad = True
+        self.assertTrue(
+            gradcheck(
+                CubicFeatureSamplingFunction.apply,
+                [ptcloud.double().cuda(), cubic_features.double().cuda(), 2],
+            )
+        )
+
+    @run_if_cuda
+    def test_neighborhood_size_3(self):
+        ptcloud = torch.rand(1, 32, 3) * 2 - 1
+        cubic_features = torch.rand(1, 2, 16, 16, 16)
+        ptcloud.requires_grad = True
+        cubic_features.requires_grad = True
+        self.assertTrue(
+            gradcheck(
+                CubicFeatureSamplingFunction.apply,
+                [ptcloud.double().cuda(), cubic_features.double().cuda(), 3],
+            )
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

torch_points_kernels/torchpoints.py

@@ -229,7 +229,9 @@ def forward(ctx, xyz1, xyz2):
     @staticmethod
     def backward(ctx, grad_dist1, grad_dist2):
         xyz1, xyz2, idx1, idx2 = ctx.saved_tensors
-        grad_xyz1, grad_xyz2 = tpcuda.chamfer_dist_grad(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2)
+        grad_xyz1, grad_xyz2 = tpcuda.chamfer_dist_grad(
+            xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2
+        )
         return grad_xyz1, grad_xyz2
 
 
@@ -260,3 +262,56 @@ def chamfer_dist(xyz1, xyz2, ignore_zeros=False):
 
     dist1, dist2 = ChamferFunction.apply(xyz1, xyz2)
     return torch.mean(dist1) + torch.mean(dist2)
+
+
+class CubicFeatureSamplingFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, ptcloud, cubic_features, neighborhood_size=1):
+        scale = cubic_features.size(2)
+        point_features, grid_pt_indexes = tpcuda.cubic_feature_sampling(
+            scale, neighborhood_size, ptcloud, cubic_features
+        )
+        ctx.save_for_backward(
+            torch.Tensor([scale]), torch.Tensor([neighborhood_size]), grid_pt_indexes
+        )
+        return point_features
+
+    @staticmethod
+    def backward(ctx, grad_point_features):
+        scale, neighborhood_size, grid_pt_indexes = ctx.saved_tensors
+        scale = int(scale.item())
+        neighborhood_size = int(neighborhood_size.item())
+        grad_point_features = grad_point_features.contiguous()
+        grad_ptcloud, grad_cubic_features = tpcuda.cubic_feature_sampling_grad(
+            scale, neighborhood_size, grad_point_features, grid_pt_indexes
+        )
+        return grad_ptcloud, grad_cubic_features, None
+
+
+def cubic_feature_sampling(ptcloud, cubic_features, neighborhood_size=1):
+    r"""
+    Sample the features of points from 3D feature maps that the point lies in.
+    Please refer to https://arxiv.org/pdf/2006.03761 for more information
+
+    Parameters
+    ----------
+    ptcloud : torch.Tensor (dtype=torch.float32)
+        (B, n_pts, 3) point clouds containing n_pts points
+    cubic_features : torch.Tensor (dtype=torch.float32)
+        (B, c, m, m, m) 3D feature maps of sizes m x m x m and c channels
+    neighborhood_size : int
+        The neighborhood cubes to sample. 
+        neighborhood_size = 1 means to sample the cube that point lies in.
+        neighborhood_size = 2 means to sample surrouding cubes (step = 1) of 
+        the cube that point lies in.
+
+    Returns
+    -------
+    dist: torch.Tensor
+        (B, n_pts, n_vertices, c), where n_vertices = (neighborhood_size * 2)^3
+    """
+    h_scale = cubic_features.size(2) / 2
+    ptcloud = ptcloud * h_scale + h_scale
+    return CubicFeatureSamplingFunction.apply(
+        ptcloud, cubic_features, neighborhood_size
+    )