Basic region growing clustering

Author: nicolas-chaulet

benchmark/region_cluster.py

@@ -0,0 +1,41 @@
+import unittest
+import torch
+import os
+import sys
+import time
+import random
+
+ROOT = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")
+sys.path.insert(0, ROOT)
+
+from torch_points_kernels.cluster import grow_proximity
+
+torch.manual_seed(0)
+
+num_points = 100000
+pos1 = torch.rand((num_points, 3))
+pos2 = torch.rand((num_points, 3)) + 2
+pos3 = torch.rand((num_points, 3)) + 4
+labels1 = torch.ones(num_points).long()
+labels2 = torch.ones(num_points).long()
+labels3 = torch.ones(num_points).long()
+pos = torch.cat([pos1, pos2, pos3], 0)
+label = torch.cat([labels1, labels2, labels3], 0)
+batch = torch.ones((3 * num_points)).long()
+cl = grow_proximity(pos, batch, radius=0.5)
+
+
+import cProfile, pstats, io
+from pstats import SortKey
+
+pr = cProfile.Profile()
+pr.enable()
+t_start = time.perf_counter()
+grow_proximity(pos, batch)
+print(time.perf_counter() - t_start)
+pr.disable()
+s = io.StringIO()
+sortby = SortKey.CUMULATIVE
+ps = pstats.Stats(pr, stream=s).sort_stats(sortby)
+ps.print_stats()
+print(s.getvalue())

test/test_cluster.py

@@ -0,0 +1,52 @@
+import unittest
+import torch
+import os
+import sys
+
+ROOT = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")
+sys.path.insert(0, ROOT)
+
+from torch_points_kernels.cluster import grow_proximity, region_grow
+
+
+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],
+            ]
+        )
+        self.batch = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1])
+        self.labels = torch.tensor([0, 0, 1, 1, 0, 1, 1, 10])
+
+    def test_simple(self):
+        clusters = grow_proximity(self.pos, self.batch, radius=2, min_cluster_size=1)
+        self.assertEqual(clusters, [[0, 1, 2], [3], [4, 5, 6], [7]])
+
+        clusters = grow_proximity(self.pos, self.batch, radius=2, min_cluster_size=3)
+        self.assertEqual(clusters, [[0, 1, 2], [4, 5, 6]])
+
+    def test_region_grow(self):
+        clusters = region_grow(
+            self.pos, self.labels, self.batch, radius=2, min_cluster_size=1
+        )
+        self.assertEqual(len(clusters[0]), 2)
+        self.assertEqual(len(clusters[1]), 3)
+        self.assertEqual(len(clusters[10]), 1)
+        torch.testing.assert_allclose(clusters[0][0], torch.tensor([0, 1]))
+        torch.testing.assert_allclose(clusters[0][1], torch.tensor([4]))
+        torch.testing.assert_allclose(clusters[1][0], torch.tensor([2]))
+        torch.testing.assert_allclose(clusters[1][1], torch.tensor([3]))
+        torch.testing.assert_allclose(clusters[1][2], torch.tensor([5, 6]))
+        torch.testing.assert_allclose(clusters[10][0], torch.tensor([7]))
+
+
+if __name__ == "__main__":
+    unittest.main()

torch_points_kernels/cluster.py

@@ -0,0 +1,87 @@
+import torch
+from .torchpoints import ball_query_partial_dense
+import numpy as np
+import numba
+
+
+@numba.jit(nopython=True)
+def _grow_proximity_core(neighbours, min_cluster_size):
+    num_points = int(neighbours.shape[0])
+    visited = np.zeros((num_points,), dtype=numba.types.bool_)
+    clusters = []
+    for i in range(num_points):
+        if visited[i]:
+            continue
+
+        cluster = []
+        queue = []
+        visited[i] = True
+        queue.append(i)
+        cluster.append(i)
+        while len(queue):
+            k = queue.pop()
+            k_neighbours = neighbours[k]
+            for nei in k_neighbours:
+                if nei.item() == -1:
+                    break
+
+                if not visited[nei]:
+                    visited[nei] = True
+                    queue.append(nei.item())
+                    cluster.append(nei.item())
+
+        if len(cluster) >= min_cluster_size:
+            clusters.append(cluster)
+
+    return clusters
+
+
+def grow_proximity(pos, batch, nsample=16, radius=0.02, min_cluster_size=32):
+    """ Grow based on proximity only"""
+    assert pos.shape[0] == batch.shape[0]
+    neighbours = (
+        ball_query_partial_dense(radius, nsample, pos, pos, batch, batch)[0]
+        .cpu()
+        .numpy()
+    )
+    return _grow_proximity_core(neighbours, min_cluster_size)
+
+
+def region_grow(
+    pos, labels, batch, ignore_labels=[], nsample=16, radius=0.02, min_cluster_size=32
+):
+    """ 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
+
+    Parameters
+    ----------
+    """
+    assert labels.dim() == 1
+    assert pos.dim() == 2
+    assert pos.shape[0] == labels.shape[0]
+
+    unique_labels = torch.unique(labels)
+    clusters = {}
+    ind = torch.arange(0, pos.shape[0])
+    for l in unique_labels:
+        if l in ignore_labels:
+            continue
+        label_mask = labels == l
+        local_ind = ind[label_mask]
+        label_clusters = grow_proximity(
+            pos[label_mask, :],
+            batch[label_mask],
+            nsample=nsample,
+            radius=radius,
+            min_cluster_size=min_cluster_size,
+        )
+        if len(label_clusters):
+            remaped_clusters = []
+            for cluster in label_clusters:
+                cluster = cluster.to(pos.device)
+                remaped_clusters.append(local_ind[cluster])
+            clusters[l.item()] = remaped_clusters
+
+    return clusters

torch_points_kernels/torchpoints.py

@@ -30,7 +30,10 @@ def furthest_point_sample(xyz, npoint):
         (B, npoint) tensor containing the set
     """
     if npoint > xyz.shape[1]:
-        raise ValueError("caanot sample %i points from an input set of %i points" % (npoint, xyz.shape[1]))
+        raise ValueError(
+            "caanot sample %i points from an input set of %i points"
+            % (npoint, xyz.shape[1])
+        )
     if xyz.is_cuda:
         return tpcuda.furthest_point_sampling(xyz, npoint)
     else:
@@ -99,9 +102,13 @@ def backward(ctx, grad_out):
         idx, weight, m = ctx.three_interpolate_for_backward
 
         if grad_out.is_cuda:
-            grad_features = tpcuda.three_interpolate_grad(grad_out.contiguous(), idx, weight, m)
+            grad_features = tpcuda.three_interpolate_grad(
+                grad_out.contiguous(), idx, weight, m
+            )
         else:
-            grad_features = tpcpu.knn_interpolate_grad(grad_out.contiguous(), idx, weight, m)
+            grad_features = tpcpu.knn_interpolate_grad(
+                grad_out.contiguous(), idx, weight, m
+            )
 
         return grad_features, None, None
 
@@ -143,17 +150,23 @@ def grouping_operation(features, idx):
     all_idx = idx.reshape(idx.shape[0], -1)
     all_idx = all_idx.unsqueeze(1).repeat(1, features.shape[1], 1)
     grouped_features = features.gather(2, all_idx)
-    return grouped_features.reshape(idx.shape[0], features.shape[1], idx.shape[1], idx.shape[2])
+    return grouped_features.reshape(
+        idx.shape[0], features.shape[1], idx.shape[1], idx.shape[2]
+    )
 
 
-def ball_query_dense(radius, nsample, xyz, new_xyz, batch_xyz=None, batch_new_xyz=None, sort=False):
+def ball_query_dense(
+    radius, nsample, xyz, new_xyz, batch_xyz=None, batch_new_xyz=None, sort=False
+):
     # type: (Any, float, int, torch.Tensor, torch.Tensor) -> torch.Tensor
     if new_xyz.is_cuda:
         if sort:
             raise NotImplementedError("CUDA version does not sort the neighbors")
         ind, dist = tpcuda.ball_query_dense(new_xyz, xyz, radius, nsample)
     else:
-        ind, dist = tpcpu.dense_ball_query(new_xyz, xyz, radius, nsample, mode=0, sorted=sort)
+        ind, dist = tpcpu.dense_ball_query(
+            new_xyz, xyz, radius, nsample, mode=0, sorted=sort
+        )
     return ind, dist
 
 
@@ -162,9 +175,13 @@ def ball_query_partial_dense(radius, nsample, x, y, batch_x, batch_y, sort=False
     if x.is_cuda:
         if sort:
             raise NotImplementedError("CUDA version does not sort the neighbors")
-        ind, dist = tpcuda.ball_query_partial_dense(x, y, batch_x, batch_y, radius, nsample)
+        ind, dist = tpcuda.ball_query_partial_dense(
+            x, y, batch_x, batch_y, radius, nsample
+        )
     else:
-        ind, dist = tpcpu.batch_ball_query(x, y, batch_x, batch_y, radius, nsample, mode=0, sorted=sort)
+        ind, dist = tpcpu.batch_ball_query(
+            x, y, batch_x, batch_y, radius, nsample, mode=0, sorted=sort
+        )
     return ind, dist
 
 
@@ -207,7 +224,9 @@ def ball_query(
         assert x.size(0) == batch_x.size(0)
         assert y.size(0) == batch_y.size(0)
         assert x.dim() == 2
-        return ball_query_partial_dense(radius, nsample, x, y, batch_x, batch_y, sort=sort)
+        return ball_query_partial_dense(
+            radius, nsample, x, y, batch_x, batch_y, sort=sort
+        )
 
     elif mode.lower() == "dense":
         if (batch_x is not None) or (batch_y is not None):