CUDA RANSAC implementation

.github/workflows/test.yml

@@ -44,6 +44,8 @@ jobs:
       run: |
         pip install .
         pytest
+      env:
+        NUMBA_ENABLE_CUDASIM: 1
 
   publish-package:
     if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')

octreelib/grid/grid.py

@@ -13,6 +13,7 @@
 )
 from octreelib.internal.point import PointCloud
 from octreelib.internal.voxel import Voxel, VoxelBase
+from octreelib.ransac.cuda_ransac import CudaRansac
 
 __all__ = ["Grid", "GridConfig"]
 
@@ -107,23 +108,127 @@ def insert_points(self, pose_number: int, points: PointCloud):
             self.__pose_voxel_coordinates[pose_number].append(target_voxel)
             self.__octrees[target_voxel].insert_points(pose_number, voxel_points)
 
-    def map_leaf_points(self, function: Callable[[PointCloud], PointCloud]):
+    def map_leaf_points(
+        self,
+        function: Callable[[PointCloud], PointCloud],
+        pose_numbers: Optional[List[int]] = None,
+    ):
         """
         Transform point cloud in each node of each octree using the function
         :param function: Transformation function PointCloud -> PointCloud. It is applied to each leaf node.
+        :param pose_numbers: List of pose numbers to map.
         """
         for voxel_coordinates in self.__octrees:
-            self.__octrees[voxel_coordinates].map_leaf_points(function)
+            self.__octrees[voxel_coordinates].map_leaf_points(function, pose_numbers)
+
+    def map_leaf_points_cuda_ransac(
+        self,
+        poses_per_batch: int = 1,
+        threshold: float = 0.01,
+        n_hypotheses: int = 1024,
+    ):
+        """
+        transform point cloud in the node using the function
+        :param poses_per_batch: Number of poses per batch.
+        :param threshold: Distance threshold.
+        :param n_hypotheses: Number of RANSAC iterations (<= 1024).
+        """
+        if threshold <= 0:
+            raise ValueError("Threshold must be positive")
+        if n_hypotheses < 1:
+            raise ValueError("Number of RANSAC hypotheses must be positive")
+        if n_hypotheses > 1024:
+            raise ValueError(
+                "Number of RANSAC hypotheses must be <= 1024 "
+                "because of the CUDA thread limit."
+            )
+
+        # processing is done in batches to avoid running out of memory
+        batches = [
+            list(
+                range(
+                    i,
+                    min(i + poses_per_batch, len(self.__pose_voxel_coordinates)),
+                )
+            )
+            for i in range(0, len(self.__pose_voxel_coordinates), poses_per_batch)
+        ]
 
-    def get_leaf_points(self, pose_number: int) -> List[Voxel]:
+        # find the maximum number of leaf voxels across all batches
+        # this is needed to initialize the random number generators on the GPU
+        max_leaf_voxels = max(
+            [
+                sum([self.n_leaves(pose_number) for pose_number in batch_pose_numbers])
+                for batch_pose_numbers in batches
+            ]
+        )
+        ransac = CudaRansac(
+            max_n_blocks=max_leaf_voxels,
+            n_hypotheses=n_hypotheses,
+            threshold=threshold,
+        )
+
+        # process each batch
+        for batch_pose_numbers in batches:
+            # `combined_point_cloud` is a concatenation of ALL point clouds
+            # `block_sizes` is a list of sizes of point clouds for each leaf node
+            # `pose_dividers` is a list of indices where combined_point_cloud is divided by pose
+            # these are used to split the combined_point_cloud into separate point clouds
+            # for each pose after the kernel is done
+            batch_point_clouds = []
+            block_sizes = []
+            pose_dividers = [0]
+            for pose_number in batch_pose_numbers:
+                pose_point_cloud = self.get_leaf_points(pose_number)
+                batch_point_clouds.append(
+                    np.vstack([v.get_points() for v in pose_point_cloud])
+                )
+                block_sizes.append(
+                    np.array(
+                        [len(v.get_points()) for v in pose_point_cloud],
+                        dtype=np.int32,
+                    )
+                )
+                pose_dividers.append(pose_dividers[-1] + block_sizes[-1].sum())
+
+            combined_point_cloud = np.vstack(batch_point_clouds)
+            block_sizes_combined = np.concatenate(block_sizes)
+            pose_dividers = np.array(pose_dividers)
+
+            # run the kernel
+            maximum_mask = ransac.evaluate(
+                combined_point_cloud,
+                block_sizes_combined,
+            )
+
+            # split the combined point cloud into separate point clouds for each pose,
+            # apply the masks from the kernel
+            # and insert them into the octrees
+
+            for i, (pose_number, block_sizes_for_pose) in enumerate(
+                zip(batch_pose_numbers, block_sizes)
+            ):
+                mask = maximum_mask[pose_dividers[i] : pose_dividers[i + 1]]
+                point_start_index = 0
+                for voxel_coordinates in self.__pose_voxel_coordinates[pose_number]:
+                    octree = self.__octrees[voxel_coordinates]
+                    n_points = octree.n_points(pose_number)
+                    octree_mask = mask[point_start_index : point_start_index + n_points]
+                    octree.apply_mask(octree_mask, pose_number)
+                    point_start_index += n_points
+
+    def get_leaf_points(self, pose_number: int, non_empty: bool = True) -> List[Voxel]:
         """
         :param pose_number: The desired pose number.
+        :param non_empty: If True, returns only non-empty voxels.
         :return: List of voxels. Each voxel is a representation of a leaf node.
         Each voxel has the same corner, edge_length and points as one of the leaf nodes.
         """
         return sum(
             [
-                self.__octrees[voxel_coordinates].get_leaf_points(pose_number)
+                self.__octrees[voxel_coordinates].get_leaf_points(
+                    non_empty, pose_number
+                )
                 for voxel_coordinates in self.__pose_voxel_coordinates[pose_number]
             ],
             [],

octreelib/grid/grid_base.py

@@ -154,6 +154,21 @@ def map_leaf_points(self, function: Callable[[PointCloud], PointCloud]):
         """
         pass
 
+    @abstractmethod
+    def map_leaf_points_cuda_ransac(
+        self,
+        n_poses_per_batch: int = 10,
+        threshold: int = 0.01,
+        n_iterations: int = 1024,
+    ):
+        """
+        transform point cloud in the node using the function
+        :param n_poses_per_batch: Number of poses per batch.
+        :param threshold: Distance threshold.
+        :param n_iterations: Number of RANSAC iterations (<= 1024).
+        """
+        pass
+
     @abstractmethod
     def get_leaf_points(self, pose_number: int) -> List[Voxel]:
         """

octreelib/octree/octree.py

@@ -48,6 +48,8 @@ def subdivide_as(self, other: "OctreeNode"):
         elif self._has_children:
             self._points = self.get_points()
             self._has_children = False
+            for child in self._children:
+                child._remove_from_cache()
             self._children = []
 
     def get_points(self) -> PointCloud:
@@ -132,6 +134,13 @@ def get_leaf_points(self) -> List[Voxel]:
             else []
         )
 
+    def apply_mask(self, mask: np.ndarray):
+        """
+        Apply mask to the point cloud in the octree node
+        :param mask: Mask to apply
+        """
+        self._points = self._points[mask]
+
     @property
     def n_leaves(self):
         """
@@ -171,14 +180,25 @@ def _generate_children(self):
         """
         child_edge_length = self.edge_length / np.float_(2)
         children_corners_offsets = itertools.product([0, child_edge_length], repeat=3)
+        self._cached_leaves.remove(self)
         return [
             OctreeNode(
                 self.corner_min + offset,
                 child_edge_length,
+                self._cached_leaves,
             )
             for internal_position, offset in enumerate(children_corners_offsets)
         ]
 
+    def _remove_from_cache(self):
+        """
+        Remove the node and its children from the cached leaves.
+        """
+        self._cached_leaves.remove(self)
+        if self._has_children:
+            for child in self._children:
+                child._remove_from_cache()
+
 
 class Octree(OctreeBase, Generic[T]):
     """
@@ -233,11 +253,25 @@ def map_leaf_points(self, function: Callable[[PointCloud], PointCloud]):
         """
         self._root.map_leaf_points(function)
 
-    def get_leaf_points(self) -> List[Voxel]:
+    def get_leaf_points(self, non_empty: bool = True) -> List[Voxel]:
         """
+        :param non_empty: If True, only non-empty leaf nodes are returned.
         :return: List of voxels where each voxel represents a leaf node with points.
         """
-        return self._root.get_leaf_points()
+        if non_empty:
+            return list(filter(lambda v: v.n_points != 0, self._cached_leaves))
+        return self._cached_leaves
+
+    def apply_mask(self, mask: np.ndarray):
+        """
+        Apply mask to the point cloud in the octree
+        :param mask: Mask to apply
+        """
+        start_index = 0
+        for leaf in filter(lambda v: v.n_points != 0, self._cached_leaves):
+            n_points = leaf.n_points
+            leaf.apply_mask(mask[start_index : start_index + n_points])
+            start_index += n_points
 
     @property
     def n_points(self):

octreelib/octree/octree_base.py

@@ -33,11 +33,20 @@ class OctreeNodeBase(Voxel, ABC):
     and are not stored in the parent node.
     """
 
-    def __init__(self, corner_min: Point, edge_length: float):
+    def __init__(
+        self,
+        corner_min: Point,
+        edge_length: float,
+        octree_cached_leaves: List["OctreeNodeBase"],
+    ):
         super().__init__(corner_min, edge_length)
         self._points: np.empty((0, 3), dtype=float)
         self._children: Optional[List["OctreeNodeBase"]] = []
         self._has_children: bool = False
+        # `OctreeNodeBase_cached_leaves` references field `OctreeBase._cached_leaves`
+        # so that nodes can modify this field in the parent OctreeBase instance
+        self._cached_leaves = octree_cached_leaves
+        self._cached_leaves.append(self)
 
     @property
     @abstractmethod
@@ -111,6 +120,14 @@ def get_points(self) -> PointCloud:
         """
         pass
 
+    @abstractmethod
+    def apply_mask(self, mask: np.ndarray):
+        """
+        Apply mask to the point cloud in the octree node
+        :param mask: Mask to apply
+        """
+        self._points = self._points[mask]
+
 
 class OctreeBase(Voxel, ABC):
     """
@@ -131,7 +148,14 @@ def __init__(
     ):
         super().__init__(corner_min, edge_length)
         self._config = octree_config
-        self._root = self._node_type(self.corner_min, self.edge_length)
+
+        # cached leaves allow for fast retrieval of leaf nodes with points
+        # skipping the stage of finding them and returning through multiple
+        # layers of recursion
+        self._cached_leaves = []
+        self._root = self._node_type(
+            self.corner_min, self.edge_length, self._cached_leaves
+        )
 
     @property
     @abstractmethod
@@ -175,8 +199,9 @@ def map_leaf_points(self, function: Callable[[PointCloud], PointCloud]):
         pass
 
     @abstractmethod
-    def get_leaf_points(self) -> List[Voxel]:
+    def get_leaf_points(self, non_empty: bool) -> List[Voxel]:
         """
+        :param non_empty: If True, only non-empty leaf nodes are returned.
         :return: List of PointClouds where each PointCloud
         represents points in a separate leaf node
         """
@@ -208,3 +233,10 @@ def get_points(self) -> PointCloud:
     @abstractmethod
     def insert_points(self, points: PointCloud):
         pass
+
+    @abstractmethod
+    def apply_mask(self, mask: np.ndarray):
+        """
+        Apply mask to the point cloud in the octree
+        :param mask: Mask to apply
+        """

octreelib/octree_manager/octree_manager.py

@@ -79,7 +79,8 @@ def map_leaf_points(
             pose_numbers = self._octrees.keys()
 
         for pose_number in pose_numbers:
-            self._octrees[pose_number].map_leaf_points(function)
+            if pose_number in self._octrees:
+                self._octrees[pose_number].map_leaf_points(function)
 
     def filter(
         self,
@@ -97,17 +98,24 @@ def filter(
         for pose_number in pose_numbers:
             self._octrees[pose_number].filter(filtering_criteria)
 
-    def get_leaf_points(self, pose_number: Optional[int] = None) -> List[Voxel]:
+    def get_leaf_points(
+        self, non_empty: bool = True, pose_number: Optional[int] = None
+    ) -> List[Voxel]:
         """
+        :param non_empty: If True, only non-empty leaf nodes are returned.
         :param pose_number: Desired pose number.
         :return: List of leaf voxels with points for this pose.
         """
         if pose_number is None:
             return sum(
-                [octree.get_leaf_points() for octree in self._octrees.values()], []
+                [
+                    octree.get_leaf_points(non_empty)
+                    for octree in self._octrees.values()
+                ],
+                [],
             )
         if pose_number in self._octrees:
-            return self._octrees[pose_number].get_leaf_points()
+            return self._octrees[pose_number].get_leaf_points(non_empty)
         return []
 
     def get_points(self, pose_number: Optional[int] = None) -> PointCloud:
@@ -161,3 +169,12 @@ def insert_points(self, pose_number: int, points: PointCloud):
             )
         self._octrees[pose_number].insert_points(points)
         self._octrees[pose_number].subdivide_as(self._scheme_octree)
+
+    def apply_mask(self, mask: np.ndarray, pose_number: int):
+        """
+        Apply mask to the point cloud in the octree
+        :param mask: Mask to apply
+        :param pose_number: Pose number
+        """
+        if pose_number in self._octrees:
+            self._octrees[pose_number].apply_mask(mask)

octreelib/ransac/__init__.py

@@ -0,0 +1,9 @@
+"""
+This module contains implementations of RANSAC algorithm.
+"""
+
+import octreelib.ransac.cuda_ransac as cuda_ransac_module
+
+from octreelib.ransac.cuda_ransac import *
+
+__all__ = cuda_ransac_module.__all__