Resolve conflicts.

hzxie · hzxie · commit eda4158153d4 · 2020-12-18T09:04:28.000+08:00
diff --git a/setup.py b/setup.py
@@ -37,7 +37,10 @@ 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,
+                },
             )
         )
 
@@ -49,7 +52,9 @@ def get_ext_modules():
             name="torch_points_kernels.points_cpu",
             sources=cpu_ext_sources,
             include_dirs=["{}/include".format(cpu_ext_src_root)],
-            extra_compile_args={"cxx": extra_compile_args,},
+            extra_compile_args={
+                "cxx": extra_compile_args,
+            },
         )
     )
     return ext_modules
@@ -81,5 +86,8 @@ 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",
+    ],
 )
diff --git a/test/speed_radius.py b/test/speed_radius.py
@@ -23,8 +23,26 @@ 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)
 
diff --git a/test/test_ballquerry.py b/test/test_ballquerry.py
@@ -112,12 +112,48 @@ def test_random_cpu(self, cuda=False):
         batch_b = torch.tensor([0 for i in range(b.shape[0] // 2)] + [1 for i in range(b.shape[0] // 2, b.shape[0])])
         R = 1
 
-        idx, dist = ball_query(R, 15, a, b, mode="PARTIAL_DENSE", batch_x=batch_a, batch_y=batch_b, sort=True,)
-        idx1, dist = ball_query(R, 15, a, b, mode="PARTIAL_DENSE", batch_x=batch_a, batch_y=batch_b, sort=True,)
+        idx, dist = ball_query(
+            R,
+            15,
+            a,
+            b,
+            mode="PARTIAL_DENSE",
+            batch_x=batch_a,
+            batch_y=batch_b,
+            sort=True,
+        )
+        idx1, dist = ball_query(
+            R,
+            15,
+            a,
+            b,
+            mode="PARTIAL_DENSE",
+            batch_x=batch_a,
+            batch_y=batch_b,
+            sort=True,
+        )
         torch.testing.assert_allclose(idx1, idx)
         with self.assertRaises(AssertionError):
-            idx, dist = ball_query(R, 15, a, b, mode="PARTIAL_DENSE", batch_x=batch_a, batch_y=batch_b, sort=False,)
-            idx1, dist = ball_query(R, 15, a, b, mode="PARTIAL_DENSE", batch_x=batch_a, batch_y=batch_b, sort=False,)
+            idx, dist = ball_query(
+                R,
+                15,
+                a,
+                b,
+                mode="PARTIAL_DENSE",
+                batch_x=batch_a,
+                batch_y=batch_b,
+                sort=False,
+            )
+            idx1, dist = ball_query(
+                R,
+                15,
+                a,
+                b,
+                mode="PARTIAL_DENSE",
+                batch_x=batch_a,
+                batch_y=batch_b,
+                sort=False,
+            )
             torch.testing.assert_allclose(idx1, idx)
 
         self.assertEqual(idx.shape[0], b.shape[0])
@@ -144,7 +180,16 @@ def test_random_gpu(self):
         ).cuda()
         R = 1
 
-        idx, dist = ball_query(R, 15, a, b, mode="PARTIAL_DENSE", batch_x=batch_a, batch_y=batch_b, sort=False,)
+        idx, dist = ball_query(
+            R,
+            15,
+            a,
+            b,
+            mode="PARTIAL_DENSE",
+            batch_x=batch_a,
+            batch_y=batch_b,
+            sort=False,
+        )
 
         # Comparison to see if we have the same result
         tree = KDTree(a.cpu().detach().numpy())
diff --git a/test/test_chamfer_dist.py b/test/test_chamfer_dist.py
@@ -12,7 +12,7 @@
 ROOT = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")
 sys.path.insert(0, ROOT)
 
-from torch_points_kernels import ChamferFunction, chamfer_dist
+from torch_points_kernels.chamfer_dist import ChamferFunction, chamfer_dist
 
 
 class TestChamferDistance(unittest.TestCase):
diff --git a/test/test_cluster.py b/test/test_cluster.py
@@ -12,7 +12,16 @@
 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])
diff --git a/test/test_grouping.py b/test/test_grouping.py
@@ -13,7 +13,11 @@ def test_simple(self):
         features = torch.tensor(
             [
                 [[0, 10, 0], [1, 11, 0], [2, 12, 0]],
-                [[100, 110, 120], [101, 111, 121], [102, 112, 122],],  # x-coordinates  # y-coordinates  # z-coordinates
+                [
+                    [100, 110, 120],
+                    [101, 111, 121],
+                    [102, 112, 122],
+                ],  # x-coordinates  # y-coordinates  # z-coordinates
             ]
         ).type(torch.float)
         idx = torch.tensor([[[1, 0], [0, 0]], [[0, 1], [1, 2]]]).type(torch.long)
@@ -38,7 +42,8 @@ def test_simple(self):
 
         if torch.cuda.is_available():
             npt.assert_array_equal(
-                grouping_operation(features.cuda(), idx.cuda()).detach().cpu().numpy(), expected,
+                grouping_operation(features.cuda(), idx.cuda()).detach().cpu().numpy(),
+                expected,
             )
 
 
diff --git a/test/test_metrics.py b/test/test_metrics.py
@@ -25,7 +25,8 @@ def test_simple(self, cuda=False):
             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]]),
+            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):
@@ -42,7 +43,14 @@ def test_batch(self, cuda=False):
             batch = batch.cuda()
         ious = instance_iou(proposed_instances, gt_instances, batch=batch)
         torch.testing.assert_allclose(
-            ious.cpu(), torch.tensor([[0.5, 0.5, 0, 0, 0], [0, 0, 0, 1, 0], [0, 0, 0, 0, 1],]),
+            ious.cpu(),
+            torch.tensor(
+                [
+                    [0.5, 0.5, 0, 0, 0],
+                    [0, 0, 0, 1, 0],
+                    [0, 0, 0, 0, 1],
+                ]
+            ),
         )
 
     @run_if_cuda
diff --git a/torch_points_kernels/chamfer_dist.py b/torch_points_kernels/chamfer_dist.py
@@ -0,0 +1,59 @@
+import torch
+
+if torch.cuda.is_available():
+    import torch_points_kernels.points_cuda as tpcuda
+
+
+class ChamferFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, xyz1, xyz2):
+        if not torch.cuda.is_available():
+            raise NotImplementedError(
+                "CPU version is not available for Chamfer Distance"
+            )
+
+        dist1, dist2, idx1, idx2 = tpcuda.chamfer_dist(xyz1, xyz2)
+        ctx.save_for_backward(xyz1, xyz2, idx1, idx2)
+
+        return dist1, dist2
+
+    @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
+        )
+        return grad_xyz1, grad_xyz2
+
+
+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
+    """
+    if len(xyz1.shape) != 3 or xyz1.size(2) != 3 or len(xyz2.shape) != 3 or xyz2.size(2) != 3:
+        raise ValueError('The input point cloud should be of size (B, n_pts, 3)')
+
+    batch_size = xyz1.size(0)
+    if batch_size == 1 and ignore_zeros:
+        non_zeros1 = torch.sum(xyz1, dim=2).ne(0)
+        non_zeros2 = torch.sum(xyz2, dim=2).ne(0)
+        xyz1 = xyz1[non_zeros1].unsqueeze(dim=0)
+        xyz2 = xyz2[non_zeros2].unsqueeze(dim=0)
+
+    dist1, dist2 = ChamferFunction.apply(xyz1, xyz2)
+    return torch.mean(dist1) + torch.mean(dist2)
+
diff --git a/torch_points_kernels/cluster.py b/torch_points_kernels/cluster.py
@@ -38,7 +38,7 @@ def _grow_proximity_core(neighbours, min_cluster_size):
 
 
 def grow_proximity(pos, batch, nsample=16, radius=0.02, min_cluster_size=32):
-    """ Grow based on proximity only
+    """Grow based on proximity only
     Neighbour search is done on device while the cluster assignement is done on cpu"""
     assert pos.shape[0] == batch.shape[0]
     neighbours = ball_query_partial_dense(radius, nsample, pos, pos, batch, batch)[0].cpu().numpy()
@@ -48,7 +48,7 @@ def grow_proximity(pos, batch, nsample=16, radius=0.02, min_cluster_size=32):
 def region_grow(
     pos, labels, batch, ignore_labels=[], nsample=16, radius=0.02, min_cluster_size=32
 ) -> List[torch.Tensor]:
-    """ Region growing clustering algorithm proposed in
+    """Region growing clustering algorithm proposed in
     PointGroup: Dual-Set Point Grouping for 3D Instance Segmentation
     https://arxiv.org/pdf/2004.01658.pdf
     for instance segmentation
@@ -93,7 +93,11 @@ def region_grow(
 
         # Cluster
         label_clusters = grow_proximity(
-            pos[label_mask, :], remaped_batch, nsample=nsample, radius=radius, min_cluster_size=min_cluster_size,
+            pos[label_mask, :],
+            remaped_batch,
+            nsample=nsample,
+            radius=radius,
+            min_cluster_size=min_cluster_size,
         )
 
         # Remap indices to original coordinates
diff --git a/torch_points_kernels/knn.py b/torch_points_kernels/knn.py
@@ -2,7 +2,7 @@
 
 
 def knn(pos_support, pos, k):
-    """ Dense knn serach
+    """Dense knn serach
     Arguments:
         pos_support - [B,N,3] support points
         pos - [B,M,3] centre of queries
diff --git a/torch_points_kernels/metrics.py b/torch_points_kernels/metrics.py
@@ -9,7 +9,12 @@
 
 @numba.jit(nopython=True, parallel=True)
 def _instance_iou_cpu(
-    instance_idx, instance_offsets, gt_instances, gt_instance_sizes, num_gt_instances: np.array, batch: np.array,
+    instance_idx,
+    instance_offsets,
+    gt_instances,
+    gt_instance_sizes,
+    num_gt_instances: np.array,
+    batch: np.array,
 ):
     num_proposed_instances = len(instance_offsets) - 1
     iou = np.zeros((num_proposed_instances, num_gt_instances.sum()))
@@ -31,9 +36,11 @@ def _instance_iou_cpu(
 
 
 def instance_iou(
-    instance_idx: List[torch.Tensor], gt_instances: torch.Tensor, batch: Optional[torch.Tensor] = None,
+    instance_idx: List[torch.Tensor],
+    gt_instances: torch.Tensor,
+    batch: Optional[torch.Tensor] = None,
 ):
-    """ Computes the IoU between each proposed instance in instance_idx and ground truth instances. Returns a
+    """Computes the IoU between each proposed instance in instance_idx and ground truth instances. Returns a
     tensor of shape [instance_idx.shape[0], num_instances] that contains the iou between the proposed instances and all gt instances
     Instance label 0 is reserved for non instance points
 
diff --git a/torch_points_kernels/torchpoints.py b/torch_points_kernels/torchpoints.py
@@ -217,46 +217,3 @@ def ball_query(
     else:
         raise Exception("unrecognized mode {}".format(mode))
 
-
-class ChamferFunction(Function):
-    @staticmethod
-    def forward(ctx, xyz1, xyz2):
-        dist1, dist2, idx1, idx2 = tpcuda.chamfer_dist(xyz1, xyz2)
-        ctx.save_for_backward(xyz1, xyz2, idx1, idx2)
-
-        return dist1, dist2
-
-    @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)
-        return grad_xyz1, grad_xyz2
-
-
-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)
-        non_zeros2 = torch.sum(xyz2, dim=2).ne(0)
-        xyz1 = xyz1[non_zeros1].unsqueeze(dim=0)
-        xyz2 = xyz2[non_zeros2].unsqueeze(dim=0)
-
-    dist1, dist2 = ChamferFunction.apply(xyz1, xyz2)
-    return torch.mean(dist1) + torch.mean(dist2)

Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,11 @@ def test_simple(self):`
`13`	`13`	`features = torch.tensor(`
`14`	`14`	`[`
`15`	`15`	`[[0, 10, 0], [1, 11, 0], [2, 12, 0]],`
`16`		`- [[100, 110, 120], [101, 111, 121], [102, 112, 122],], # x-coordinates # y-coordinates # z-coordinates`
	`16`	`+ [`
	`17`	`+ [100, 110, 120],`
	`18`	`+ [101, 111, 121],`
	`19`	`+ [102, 112, 122],`
	`20`	`+ ], # x-coordinates # y-coordinates # z-coordinates`
`17`	`21`	`]`
`18`	`22`	`).type(torch.float)`
`19`	`23`	`idx = torch.tensor([[[1, 0], [0, 0]], [[0, 1], [1, 2]]]).type(torch.long)`
`@@ -38,7 +42,8 @@ def test_simple(self):`
`38`	`42`
`39`	`43`	`if torch.cuda.is_available():`
`40`	`44`	`npt.assert_array_equal(`
`41`		`- grouping_operation(features.cuda(), idx.cuda()).detach().cpu().numpy(), expected,`
	`45`	`+ grouping_operation(features.cuda(), idx.cuda()).detach().cpu().numpy(),`
	`46`	`+ expected,`
`42`	`47`	`)`
`43`	`48`
`44`	`49`