Create more test cases for the forward function.

Author: hzxie

test/test_chamfer_dist.py

@@ -1,24 +1,51 @@
+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 ChamferFunction
+from torch_points_kernels import ChamferFunction, chamfer_dist
 
 
 class TestChamferDistance(unittest.TestCase):
-    def test_chamfer_dist(self):
+    @run_if_cuda
+    def test_chamfer_dist_grad(self):
         x = torch.rand(4, 64, 3).double()
         y = torch.rand(4, 128, 3).double()
         x.requires_grad = True
         y.requires_grad = True
         test = gradcheck(ChamferFunction.apply, [x.cuda(), y.cuda()])
 
+    @run_if_cuda
+    def test_chamfer_dist(self):
+        xyz1 = torch.from_numpy(np.array([[
+            [0, 0, 0],
+            [1, 1, 1],
+            [2, 0, 1]
+        ]])).float()
+        xyz2 = torch.from_numpy(np.array([[[1, 0, 0], [1, 2, 1]]])).float()
+        dist = chamfer_dist(xyz1.cuda(), xyz2.cuda())
+        self.assertAlmostEqual(dist.item(), 2.333333, places=5)
+
+    @run_if_cuda
+    def test_chamfer_dist_ignore_zeros(self):
+        xyz1 = torch.from_numpy(np.array([[
+            [0, 0, 0],
+            [1, 1, 1],
+            [2, 0, 1]
+        ]])).float()
+        xyz2 = torch.from_numpy(np.array([[[1, 0, 0], [1, 2, 1]]])).float()
+        dist = chamfer_dist(xyz1.cuda(), xyz2.cuda(), True)
+        self.assertAlmostEqual(dist.item(), 3.0, places=5)
+
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

torch_points_kernels/torchpoints.py

@@ -252,7 +252,24 @@ def backward(ctx, grad_dist1, grad_dist2):
         return grad_xyz1, grad_xyz2
 
 
-def chamfer_dist(self, xyz1, xyz2, ignore_zeros=False):
+def chamfer_dist(xyz1, xyz2, ignore_zeros=False):
+    r"""
+    Calcuates the distance between B pairs of point clouds
+
+    Parameters
+    ----------
+    xyz1 : torch.Tensor (dtype=torch.float32)
+        (B, n1, 3) B point clouds containing n1 points
+    xyz2 : torch.Tensor (dtype=torch.float32)
+        (B, n2, 3) B point clouds containing n2 points
+    ignore_zeros : bool
+        ignore the point whose coordinate is (0, 0, 0) or not
+
+    Returns
+    -------
+    dist: torch.Tensor
+        (B, ): the distances between B pairs of point clouds
+    """
     batch_size = xyz1.size(0)
     if batch_size == 1 and ignore_zeros:
         non_zeros1 = torch.sum(xyz1, dim=2).ne(0)