Cleanup

Author: nicolas-chaulet

README.md

@@ -17,6 +17,12 @@ import torch
 import torch_points.points_cuda
 ```
 
+## Build and test
+```
+python setup.py build_ext --inplace
+python -m unittest
+```
+
 ## Projects using those kernels.
 
 [```Pytorch Point Cloud Benchmark```](https://github.com/nicolas-chaulet/deeppointcloud-benchmarks) by 

test/test_ballquerry.py

@@ -20,9 +20,7 @@ def test_simple_cpu(self):
     def test_cpu_gpu_equality(self):
         a = torch.randn(5, 1000, 3)
         npt.assert_array_equal(ball_query_dense(0.1, 17, a, a).detach().numpy(),
-                               ball_query_dense(0.1, 17, a.cuda(), a.cuda()).detach().numpy())
-
-
+                               ball_query_dense(0.1, 17, a.cuda(), a.cuda()).cpu().detach().numpy())
 
 
 if __name__ == "__main__":

test/test_grouping.py

@@ -4,6 +4,7 @@
 import numpy.testing as npt
 from torch_points import grouping_operation
 
+
 class TestGroup(unittest.TestCase):
 
     # input: points(b, c, n) idx(b, npoints, nsample)
@@ -12,29 +13,29 @@ def test_simple(self):
         features = torch.tensor([
             [[0, 10, 0], [1, 11, 0], [2, 12, 0]],
             [
-                [100, 110, 120], # x-coordinates
-                [101, 111, 121], # y-coordinates
-                [102, 112, 122], # z-coordinates
+                [100, 110, 120],  # x-coordinates
+                [101, 111, 121],  # y-coordinates
+                [102, 112, 122],  # z-coordinates
             ]
-        ])
+        ]).type(torch.float)
         idx = torch.tensor([
             [[1, 0], [0, 0]],
             [[0, 1], [1, 2]]
-        ])
+        ]).type(torch.int)
 
         expected = np.array([
             [
-                [[10, 0], [0, 0]], 
-                [[11, 1], [1, 1]], 
+                [[10, 0], [0, 0]],
+                [[11, 1], [1, 1]],
                 [[12, 2], [2, 2]]
             ],
-            [ # 2nd batch
-                [ # x-coordinates
-                    [100, 110], #x-coordinates of samples for point 0
-                    [110, 120], #x-coordinates of samples for point 1
-                ], 
-                [[101, 111], [111, 121]], # y-coordinates
-                [[102, 112], [112, 122]], # z-coordinates
+            [  # 2nd batch
+                [  # x-coordinates
+                    [100, 110],  # x-coordinates of samples for point 0
+                    [110, 120],  # x-coordinates of samples for point 1
+                ],
+                [[101, 111], [111, 121]],  # y-coordinates
+                [[102, 112], [112, 122]],  # z-coordinates
             ]
         ])
 
@@ -45,10 +46,10 @@ def test_simple(self):
         if torch.cuda.is_available():
             npt.assert_array_equal(
                 grouping_operation(
-                    features.cuda(), 
+                    features.cuda(),
                     idx.cuda()
-            ).detach().cpu().numpy(), expected)
+                ).detach().cpu().numpy(), expected)
 
 
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()

torch_points/torchpoints.py

@@ -77,12 +77,12 @@ def gather_operation(features, idx):
            (B, C, N) tensor
 
        idx : torch.Tensor
-           (B, npoint, nsample) tensor of the features to gather
+           (B, npoint) tensor of the features to gather
 
        Returns
        -------
        torch.Tensor
-           (B, C, npoint, nsample) tensor
+           (B, C, npoint) tensor
        """
     return GatherOperation.apply(features, idx)