Upload learn_region_grow/growing.py

learn_region_grow/growing.py

@@ -0,0 +1,224 @@
+"""
+Inference engine: learned region growing on a preprocessed point cloud.
+
+Given a trained LrgNet model and a point cloud (with 13-channel features),
+this module performs class-agnostic instance segmentation by growing regions
+from seed points sorted by curvature.
+"""
+
+import numpy as np
+import torch
+from typing import Tuple, List, Optional
+from .utils import sample_or_pad, center_features
+
+
+class RegionGrower:
+    """
+    Wraps a trained LrgNet and runs iterative region growing on a point cloud.
+
+    Parameters
+    ----------
+    model : LrgNet
+        Trained PyTorch model.
+    device : torch.device
+    num_inlier : int
+        Max points sampled from current region for network input (default 512).
+    num_neighbor : int
+        Max points sampled from boundary candidates (default 512).
+    add_threshold : float
+        Confidence threshold for adding a neighbor (default 0.5).
+    remove_threshold : float
+        Confidence threshold for removing an inlier (default 0.5).
+    cluster_threshold : int
+        Minimum region size to keep (default 10).
+    max_steps : int
+        Maximum growing steps per seed (default 100).
+    stuck_patience : int
+        Stop if no updates for this many consecutive steps (default 2).
+    stochastic : bool
+        If True, use random sampling weighted by confidence (paper default).
+        If False, deterministic top-k thresholding.
+    """
+
+    def __init__(self, model, device,
+                 num_inlier: int = 512,
+                 num_neighbor: int = 512,
+                 add_threshold: float = 0.5,
+                 remove_threshold: float = 0.5,
+                 cluster_threshold: int = 10,
+                 max_steps: int = 100,
+                 stuck_patience: int = 2,
+                 stochastic: bool = True):
+        self.model = model.to(device).eval()
+        self.device = device
+        self.num_inlier = num_inlier
+        self.num_neighbor = num_neighbor
+        self.add_threshold = add_threshold
+        self.remove_threshold = remove_threshold
+        self.cluster_threshold = cluster_threshold
+        self.max_steps = max_steps
+        self.stuck_patience = stuck_patience
+        self.stochastic = stochastic
+
+    @torch.no_grad()
+    def grow(self, xyz: np.ndarray, features: np.ndarray,
+             voxel_map: dict, resolution: float = 0.1) -> np.ndarray:
+        """
+        Run learned region growing on the whole point cloud.
+
+        Parameters
+        ----------
+        xyz : np.ndarray, shape (N, 3)
+            Equalized point coordinates.
+        features : np.ndarray, shape (N, 13)
+            Precomputed feature vectors.
+        voxel_map : dict
+            Voxel key -> point index mapping (from voxel_equalize).
+        resolution : float
+            Voxel grid resolution.
+
+        Returns
+        -------
+        labels : np.ndarray, shape (N,), int
+            Instance label per point. -1 = unlabeled.
+        """
+        n = len(xyz)
+        labels = np.full(n, -1, dtype=np.int32)
+        visited = np.zeros(n, dtype=bool)
+
+        # Sort seeds by curvature (channel 12), ascending -> flat surfaces first
+        seed_order = np.argsort(features[:, 12])
+
+        instance_id = 0
+        for seed_idx in seed_order:
+            if visited[seed_idx]:
+                continue
+
+            region, changed = self._grow_seed(
+                seed_idx, xyz, features, voxel_map, resolution, visited
+            )
+
+            if len(region) >= self.cluster_threshold:
+                labels[list(region)] = instance_id
+                instance_id += 1
+            visited[list(region)] = True
+
+        # Fill unlabeled points by nearest-neighbor assignment
+        unlabeled = np.where(labels == -1)[0]
+        if len(unlabeled) > 0:
+            labeled = np.where(labels >= 0)[0]
+            if len(labeled) > 0:
+                # Brute-force nearest labeled neighbor
+                from scipy.spatial import cKDTree
+                tree = cKDTree(xyz[labeled])
+                _, nn_idx = tree.query(xyz[unlabeled])
+                labels[unlabeled] = labels[labeled[nn_idx]]
+
+        return labels
+
+    def _grow_seed(self, seed_idx: int, xyz: np.ndarray, features: np.ndarray,
+                   voxel_map: dict, resolution: float, visited: np.ndarray) -> Tuple[set, bool]:
+        """
+        Grow a single region from a seed point.
+
+        Returns
+        -------
+        region : set of point indices
+        changed : bool
+            Whether the region ever grew (used by caller for bookkeeping).
+        """
+        region = {seed_idx}
+        stuck_counter = 0
+        changed = False
+
+        for step in range(self.max_steps):
+            # Find candidate neighbors on boundary
+            neighbors = self._find_boundary_neighbors(region, xyz, voxel_map, resolution)
+            # Exclude already visited seeds from other instances
+            neighbors = [p for p in neighbors if not visited[p] or p in region]
+
+            if len(neighbors) == 0:
+                break
+
+            # Sample / pad
+            inlier_pts = features[np.array(list(region), dtype=np.int64)]
+            neighbor_pts = features[np.array(neighbors, dtype=np.int64)]
+
+            inlier_s = sample_or_pad(inlier_pts, self.num_inlier)
+            neighbor_s = sample_or_pad(neighbor_pts, self.num_neighbor)
+
+            inlier_c, neighbor_c = center_features(inlier_s, neighbor_s)
+
+            # To tensor: (B, C, N)
+            inl_t = torch.from_numpy(inlier_c.T).unsqueeze(0).float().to(self.device)
+            nbr_t = torch.from_numpy(neighbor_c.T).unsqueeze(0).float().to(self.device)
+
+            add_logits, remove_logits = self.model(inl_t, nbr_t)
+            add_prob = torch.sigmoid(add_logits).cpu().numpy().squeeze()  # (Nn,)
+            remove_prob = torch.sigmoid(remove_logits).cpu().numpy().squeeze()  # (Ni,)
+
+            # Map back to actual points (first len(neighbors) entries correspond to real points)
+            n_real = len(neighbors)
+            i_real = len(region)
+
+            add_prob = add_prob[:n_real]
+            remove_prob = remove_prob[:i_real]
+
+            if self.stochastic:
+                add_mask = np.random.rand(n_real) < add_prob
+                remove_mask = np.random.rand(i_real) < remove_prob
+            else:
+                add_mask = add_prob > self.add_threshold
+                remove_mask = remove_prob > self.remove_threshold
+
+            to_add = set(np.array(neighbors)[add_mask].tolist())
+            region_list = list(region)
+            to_remove = set(np.array(region_list)[remove_mask].tolist())
+
+            new_region = (region | to_add) - to_remove
+            if new_region == region:
+                stuck_counter += 1
+                if stuck_counter >= self.stuck_patience:
+                    break
+            else:
+                stuck_counter = 0
+                changed = True
+                region = new_region
+
+        return region, changed
+
+    def _find_boundary_neighbors(self, region: set, xyz: np.ndarray,
+                                 voxel_map: dict, resolution: float) -> List[int]:
+        """
+        Find points in voxels adjacent to the current region's bounding box
+        that are not yet inside the region.
+
+        This is a 6-connected voxel search: look at all voxels that share a face
+        with any voxel occupied by the region, and return their point indices.
+        """
+        # Collect occupied voxels of the region
+        region_pts = np.array(list(region), dtype=np.int64)
+        region_voxels = set()
+        for idx in region_pts:
+            v = tuple(np.round(xyz[idx] / resolution).astype(int))
+            region_voxels.add(v)
+
+        # Find adjacent voxels
+        adjacent_voxels = set()
+        for v in region_voxels:
+            for dx, dy, dz in [(-1,0,0), (1,0,0), (0,-1,0), (0,1,0), (0,0,-1), (0,0,1)]:
+                adjacent_voxels.add((v[0]+dx, v[1]+dy, v[2]+dz))
+
+        candidates = []
+        for v in adjacent_voxels:
+            if v in voxel_map and voxel_map[v] not in region:
+                candidates.append(voxel_map[v])
+
+        # Deduplicate
+        seen = set()
+        unique = []
+        for c in candidates:
+            if c not in seen:
+                seen.add(c)
+                unique.append(c)
+        return unique