Merge pull request #40 from nicolas-chaulet/region_growing

Basic region growing clustering

Author: nicolas-chaulet

.devcontainer/Dockerfile

@@ -29,5 +29,5 @@ RUN apt-get update \
     && rm -rf /var/lib/apt/lists/*
 
 RUN pip3 install -U pip
-RUN pip3 install torch numpy scikit-learn flake8 setuptools
+RUN pip3 install torch numpy scikit-learn flake8 setuptools numba
 RUN pip3 install torch_cluster torch_sparse torch_scatter torch_geometric

.github/workflows/deploy.yaml

@@ -17,7 +17,7 @@ jobs:
             - name: Install dependencies
               run: |
                   python -m pip install --upgrade pip
-                  pip install torch numpy scikit-learn flake8 setuptools wheel twine
+                  pip install torch numpy scikit-learn flake8 setuptools wheel twine numba
             - name: Build package
               run: |
                   python setup.py build_ext --inplace

.github/workflows/tests.yaml

@@ -22,14 +22,14 @@ jobs:
         - name: Install dependencies
           run: |
               python -m pip install --upgrade pip
-              pip install numpy scikit-learn flake8 setuptools
+              pip install numpy scikit-learn flake8 setuptools numba
               
         - name: Install torch windows + linux
           if: ${{matrix.os != 'macos-latest'}}
           run: pip install torch==1.5.0+cpu torchvision==0.6.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
         - name: Install torch macos
           if: ${{matrix.os == 'macos-latest'}}
-          run: pip install torch
+          run: pip install torch 
 
         - name: Build package
           run: |

CHANGELOG.md

@@ -1,4 +1,12 @@
-# UNRELEASED
+# 0.6.5
+
+## Additions
+- Clustering algorithm for [PointGroup](https://arxiv.org/pdf/2004.01658.pdf)
+
+## Change
+- Force no ninja for the compilation
+
+# 0.6.4
 
 ## Bug fix
 - CPU version works for MacOS

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())

setup.py

@@ -60,14 +60,19 @@ def get_ext_modules():
     return ext_modules
 
 
+class CustomBuildExtension(BuildExtension):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, no_python_abi_suffix=True, use_ninja=False, **kwargs)
+
+
 def get_cmdclass():
-    return {"build_ext": BuildExtension}
+    return {"build_ext": CustomBuildExtension}
 
 
-requirements = ["torch>=1.1.0"]
+requirements = ["torch>=1.1.0", "numba"]
 
 url = "https://github.com/nicolas-chaulet/torch-points-kernels"
-__version__ = "0.6.4"
+__version__ = "0.6.5"
 setup(
     name="torch-points-kernels",
     version=__version__,

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/__init__.py

@@ -1,4 +1,13 @@
 from .torchpoints import *
 from .knn import knn
+from .cluster import region_grow
 
-__all__ = ["ball_query", "furthest_point_sample", "grouping_operation", "three_interpolate", "three_nn", "knn"]
+__all__ = [
+    "ball_query",
+    "furthest_point_sample",
+    "grouping_operation",
+    "three_interpolate",
+    "three_nn",
+    "knn",
+    "region_grow",
+]

torch_points_kernels/cluster.py

@@ -0,0 +1,104 @@
+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
+    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()
+    )
+    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
+    ----------
+    pos: torch.Tensor [N, 3]
+        Location of the points
+    labels: torch.Tensor [N,]
+        labels of each point
+    ignore_labels:
+        Labels that should be ignored, no region growing will be performed on those
+    nsample:
+        maximum number of neighbours to consider
+    radius:
+        radius for the neighbour search
+    min_cluster_size:
+        Number of points above which a cluster is considered valid
+    """
+    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
+
+        # Build clusters for a given label (ignore other points)
+        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,
+        )
+
+        # Remap indices to original coordinates
+        if len(label_clusters):
+            remaped_clusters = []
+            for cluster in label_clusters:
+                cluster = torch.tensor(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):