Merge pull request #48 from hzxie/master

Format all code with clang-format and black.

Author: nicolas-chaulet

cuda/src/ball_query.cpp

@@ -6,10 +6,11 @@ void query_ball_point_kernel_dense_wrapper(int b, int n, int m, float radius, in
                                            const float* new_xyz, const float* xyz, int64_t* idx,
                                            float* dist_out);
 
-void query_ball_point_kernel_partial_wrapper(int64_t batch_size, int size_x, int size_y, float radius,
-                                             int nsample, const float* x, const float* y,
-                                             const int64_t* batch_x, const int64_t* batch_y,
-                                             int64_t* idx_out, float* dist_out);
+void query_ball_point_kernel_partial_wrapper(int64_t batch_size, int size_x, int size_y,
+                                             float radius, int nsample, const float* x,
+                                             const float* y, const int64_t* batch_x,
+                                             const int64_t* batch_y, int64_t* idx_out,
+                                             float* dist_out);
 
 std::pair<at::Tensor, at::Tensor> ball_query_dense(at::Tensor new_xyz, at::Tensor xyz,
                                                    const float radius, const int nsample)
@@ -71,10 +72,10 @@ std::pair<at::Tensor, at::Tensor> ball_query_partial_dense(at::Tensor x, at::Ten
     batch_y = degree(batch_y, batch_size);
     batch_y = at::cat({at::zeros(1, batch_y.options()), batch_y.cumsum(0)}, 0);
 
-    query_ball_point_kernel_partial_wrapper(batch_size, x.size(0), y.size(0), radius, nsample,
-                                            x.DATA_PTR<float>(), y.DATA_PTR<float>(),
-                                            batch_x.DATA_PTR<int64_t>(), batch_y.DATA_PTR<int64_t>(),
-                                            idx.DATA_PTR<int64_t>(), dist.DATA_PTR<float>());
+    query_ball_point_kernel_partial_wrapper(
+        batch_size, x.size(0), y.size(0), radius, nsample, x.DATA_PTR<float>(), y.DATA_PTR<float>(),
+        batch_x.DATA_PTR<int64_t>(), batch_y.DATA_PTR<int64_t>(), idx.DATA_PTR<int64_t>(),
+        dist.DATA_PTR<float>());
 
     return std::make_pair(idx, dist);
 }

cuda/src/ball_query_gpu.cu

@@ -9,7 +9,7 @@
 __global__ void query_ball_point_kernel_dense(int b, int n, int m, float radius, int nsample,
                                               const float* __restrict__ new_xyz,
                                               const float* __restrict__ xyz,
-                                              int64_t* __restrict__ idx_out, 
+                                              int64_t* __restrict__ idx_out,
                                               float* __restrict__ dist_out)
 {
     int batch_index = blockIdx.x;
@@ -51,10 +51,13 @@ __global__ void query_ball_point_kernel_dense(int b, int n, int m, float radius,
     }
 }
 
-__global__ void query_ball_point_kernel_partial_dense(
-    int size_x, int size_y, float radius, int nsample, const float* __restrict__ x,
-    const float* __restrict__ y, const int64_t* __restrict__ batch_x, const int64_t* __restrict__ batch_y,
-    int64_t* __restrict__ idx_out, float* __restrict__ dist_out)
+__global__ void query_ball_point_kernel_partial_dense(int size_x, int size_y, float radius,
+                                                      int nsample, const float* __restrict__ x,
+                                                      const float* __restrict__ y,
+                                                      const int64_t* __restrict__ batch_x,
+                                                      const int64_t* __restrict__ batch_y,
+                                                      int64_t* __restrict__ idx_out,
+                                                      float* __restrict__ dist_out)
 {
     // taken from
     // https://github.com/rusty1s/pytorch_cluster/blob/master/cuda/radius_kernel.cu
@@ -67,7 +70,7 @@ __global__ void query_ball_point_kernel_partial_dense(
     const ptrdiff_t end_idx_y = batch_y[batch_idx + 1];
     float radius2 = radius * radius;
 
-    for (ptrdiff_t n_y = start_idx_y +  threadIdx.x; n_y < end_idx_y; n_y += blockDim.x)
+    for (ptrdiff_t n_y = start_idx_y + threadIdx.x; n_y < end_idx_y; n_y += blockDim.x)
     {
         int64_t count = 0;
         for (ptrdiff_t n_x = start_idx_x; n_x < end_idx_x; n_x++)
@@ -92,19 +95,21 @@ __global__ void query_ball_point_kernel_partial_dense(
 }
 
 void query_ball_point_kernel_dense_wrapper(int b, int n, int m, float radius, int nsample,
-                                           const float* new_xyz, const float* xyz, int64_t* idx,float* dist_out)
+                                           const float* new_xyz, const float* xyz, int64_t* idx,
+                                           float* dist_out)
 {
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
     query_ball_point_kernel_dense<<<b, opt_n_threads(m), 0, stream>>>(b, n, m, radius, nsample,
-                                                                      new_xyz, xyz, idx,dist_out);
+                                                                      new_xyz, xyz, idx, dist_out);
 
     CUDA_CHECK_ERRORS();
 }
 
-void query_ball_point_kernel_partial_wrapper(int64_t batch_size, int size_x, int size_y, float radius,
-                                             int nsample, const float* x, const float* y,
-                                             const int64_t* batch_x, const int64_t* batch_y,
-                                             int64_t* idx_out, float* dist_out)
+void query_ball_point_kernel_partial_wrapper(int64_t batch_size, int size_x, int size_y,
+                                             float radius, int nsample, const float* x,
+                                             const float* y, const int64_t* batch_x,
+                                             const int64_t* batch_y, int64_t* idx_out,
+                                             float* dist_out)
 {
     query_ball_point_kernel_partial_dense<<<batch_size, TOTAL_THREADS_SPARSE>>>(
         size_x, size_y, radius, nsample, x, y, batch_x, batch_y, idx_out, dist_out);

cuda/src/metrics.cpp

@@ -3,8 +3,9 @@
 #include "utils.h"
 
 void instance_iou_kernel_wrapper(int64_t total_gt_instances, int64_t max_gt_instances,
-                                 const int64_t* nInstance, int nProposal, const int64_t* proposals_idx,
-                                 const int64_t* proposals_offset, const int64_t* instance_labels,
+                                 const int64_t* nInstance, int nProposal,
+                                 const int64_t* proposals_idx, const int64_t* proposals_offset,
+                                 const int64_t* instance_labels,
                                  const int64_t* offset_num_gt_instances, const int64_t* batch,
                                  const int64_t* instance_pointnum, float* proposals_iou);
 
@@ -41,9 +42,10 @@ at::Tensor instance_iou_cuda(at::Tensor instance_idx, at::Tensor instance_offset
         at::cat({at::zeros(1, num_gt_instances.options()), num_gt_instances.cumsum(0)}, 0);
     instance_iou_kernel_wrapper(
         total_gt_instances[0], max_gt_instances[0], num_gt_instances.DATA_PTR<int64_t>(),
-        num_proposed_instances, instance_idx.DATA_PTR<int64_t>(), instance_offsets.DATA_PTR<int64_t>(),
-        gt_instances.DATA_PTR<int64_t>(), offset_num_gt_instances.DATA_PTR<int64_t>(),
-        batch.DATA_PTR<int64_t>(), gt_instance_sizes.DATA_PTR<int64_t>(), output.DATA_PTR<float>());
+        num_proposed_instances, instance_idx.DATA_PTR<int64_t>(),
+        instance_offsets.DATA_PTR<int64_t>(), gt_instances.DATA_PTR<int64_t>(),
+        offset_num_gt_instances.DATA_PTR<int64_t>(), batch.DATA_PTR<int64_t>(),
+        gt_instance_sizes.DATA_PTR<int64_t>(), output.DATA_PTR<float>());
 
     return output;
 }

cuda/src/metrics_gpu.cu

@@ -9,9 +9,9 @@
 __global__ void instance_iou_cuda_kernel(
     int64_t total_gt_instances, const int64_t* __restrict__ nInstance, int nProposal,
     const int64_t* __restrict__ proposals_idx, const int64_t* __restrict__ proposals_offset,
-    const int64_t* __restrict__ instance_labels, const int64_t* __restrict__ offset_num_gt_instances,
-    const int64_t* __restrict__ batch, const int64_t* __restrict__ instance_pointnum,
-    float* proposals_iou)
+    const int64_t* __restrict__ instance_labels,
+    const int64_t* __restrict__ offset_num_gt_instances, const int64_t* __restrict__ batch,
+    const int64_t* __restrict__ instance_pointnum, float* proposals_iou)
 {
     for (int proposal_id = blockIdx.x; proposal_id < nProposal; proposal_id += gridDim.x)
     {
@@ -48,8 +48,9 @@ __global__ void instance_iou_cuda_kernel(
 // input: instance_pointnum (total_nInst), int
 // output: proposals_iou (nProposal, total_nInst), float
 void instance_iou_kernel_wrapper(int64_t total_gt_instances, int64_t max_gt_instances,
-                                 const int64_t* nInstance, int nProposal, const int64_t* proposals_idx,
-                                 const int64_t* proposals_offset, const int64_t* instance_labels,
+                                 const int64_t* nInstance, int nProposal,
+                                 const int64_t* proposals_idx, const int64_t* proposals_offset,
+                                 const int64_t* instance_labels,
                                  const int64_t* offset_num_gt_instances, const int64_t* batch,
                                  const int64_t* instance_pointnum, float* proposals_iou)
 {

setup.py

@@ -28,9 +28,7 @@ def get_ext_modules():
         extra_compile_args += ["-DVERSION_GE_1_3"]
 
     ext_src_root = "cuda"
-    ext_sources = glob.glob("{}/src/*.cpp".format(ext_src_root)) + glob.glob(
-        "{}/src/*.cu".format(ext_src_root)
-    )
+    ext_sources = glob.glob("{}/src/*.cpp".format(ext_src_root)) + glob.glob("{}/src/*.cu".format(ext_src_root))
 
     ext_modules = []
     if CUDA_HOME:
@@ -39,10 +37,7 @@ def get_ext_modules():
                 name="torch_points_kernels.points_cuda",
                 sources=ext_sources,
                 include_dirs=["{}/include".format(ext_src_root)],
-                extra_compile_args={
-                    "cxx": extra_compile_args,
-                    "nvcc": extra_compile_args,
-                },
+                extra_compile_args={"cxx": extra_compile_args, "nvcc": extra_compile_args,},
             )
         )
 
@@ -86,8 +81,5 @@ def get_cmdclass():
     cmdclass=get_cmdclass(),
     long_description=long_description,
     long_description_content_type="text/markdown",
-    classifiers=[
-        "Programming Language :: Python :: 3",
-        "License :: OSI Approved :: MIT License",
-    ],
+    classifiers=["Programming Language :: Python :: 3", "License :: OSI Approved :: MIT License",],
 )

test/speed_radius.py

@@ -23,8 +23,8 @@ def test_speed(self):
         R = 1
         samples = 50
 
-        idx, dist = ball_query(R, samples, a, b, mode="PARTIAL_DENSE", batch_x=batch_a, batch_y=batch_b, sort=True)
-        idx1, dist = ball_query(R, samples, a, b, mode="PARTIAL_DENSE", batch_x=batch_a, batch_y=batch_b, sort=True)
+        idx, dist = ball_query(R, samples, a, b, mode="PARTIAL_DENSE", batch_x=batch_a, batch_y=batch_b, sort=True,)
+        idx1, dist = ball_query(R, samples, a, b, mode="PARTIAL_DENSE", batch_x=batch_a, batch_y=batch_b, sort=True,)
         print(time.time() - start)
         torch.testing.assert_allclose(idx1, idx)
 
@@ -40,5 +40,6 @@ def test_speed(self):
         #     if p >= 0 and p < len(batch_a):
         #         assert p in idx3_sk[i]
 
+
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main()

test/test_chamfer_dist.py

@@ -26,22 +26,14 @@ def test_chamfer_dist_grad(self):
 
     @run_if_cuda
     def test_chamfer_dist(self):
-        xyz1 = torch.from_numpy(np.array([[
-            [0, 0, 0],
-            [1, 1, 1],
-            [2, 0, 1]
-        ]])).float()
+        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()
+        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)

test/test_cluster.py

@@ -12,16 +12,7 @@
 class TestGrow(unittest.TestCase):
     def setUp(self):
         self.pos = torch.tensor(
-            [
-                [0, 0, 0],
-                [1, 0, 0],
-                [2, 0, 0],
-                [10, 0, 0],
-                [0, 0, 0],
-                [1, 0, 0],
-                [2, 0, 0],
-                [10, 0, 0],
-            ]
+            [[0, 0, 0], [1, 0, 0], [2, 0, 0], [10, 0, 0], [0, 0, 0], [1, 0, 0], [2, 0, 0], [10, 0, 0],]
         )
         self.batch = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1])
         self.labels = torch.tensor([0, 0, 1, 1, 0, 1, 1, 10])
@@ -34,9 +25,7 @@ def test_simple(self):
         self.assertEqual(clusters, [[0, 1, 2], [4, 5, 6]])
 
     def test_region_grow(self):
-        cluster_idx = region_grow(
-            self.pos, self.labels, self.batch, radius=2, min_cluster_size=1
-        )
+        cluster_idx = region_grow(self.pos, self.labels, self.batch, radius=2, min_cluster_size=1)
         self.assertEqual(len(cluster_idx), 6)
         torch.testing.assert_allclose(cluster_idx[0], torch.tensor([0, 1]))
         torch.testing.assert_allclose(cluster_idx[1], torch.tensor([4]))

test/test_grouping.py

@@ -37,7 +37,9 @@ def test_simple(self):
         npt.assert_array_equal(expected, cpu_output)
 
         if torch.cuda.is_available():
-            npt.assert_array_equal(grouping_operation(features.cuda(), idx.cuda()).detach().cpu().numpy(), expected)
+            npt.assert_array_equal(
+                grouping_operation(features.cuda(), idx.cuda()).detach().cpu().numpy(), expected,
+            )
 
 
 if __name__ == "__main__":

test/test_metrics.py

@@ -24,7 +24,9 @@ def test_simple(self, cuda=False):
             proposed_instances = [c.cuda() for c in proposed_instances]
             gt_instances = gt_instances.cuda()
         ious = instance_iou(proposed_instances, gt_instances)
-        torch.testing.assert_allclose(ious.cpu(), torch.tensor([[1, 0, 0], [0, 2 / 3.0, 0], [0, 1.0 / 4.0, 1.0 / 2.0]]))
+        torch.testing.assert_allclose(
+            ious.cpu(), torch.tensor([[1, 0, 0], [0, 2 / 3.0, 0], [0, 1.0 / 4.0, 1.0 / 2.0]]),
+        )
 
     def test_batch(self, cuda=False):
         gt_instances = torch.tensor([1, 2, 1, 2, 2, 3, 0])