Upload learn_region_grow/preprocess.py

learn_region_grow/preprocess.py

@@ -0,0 +1,169 @@
+"""Preprocessing: voxel equalization, normal / curvature estimation, feature vector."""
+
+import numpy as np
+from typing import Tuple, Optional
+
+
+def voxel_equalize(xyz: np.ndarray, resolution: float = 0.1) -> Tuple[np.ndarray, np.ndarray, dict]:
+    """
+    Voxelize a point cloud at given resolution, keeping one representative per voxel.
+
+    This step is critical: it removes density bias so that highly sampled regions
+    (e.g. close to a scanner) do not dominate the neighborhood queries later.
+
+    Parameters
+    ----------
+    xyz : np.ndarray, shape (N, 3)
+        Point coordinates.
+    resolution : float
+        Voxel grid size in the same unit as xyz (default 0.1 m).
+
+    Returns
+    -------
+    eq_xyz : np.ndarray, shape (M, 3)
+        Equalized point coordinates (M <= N).
+    eq_idx : np.ndarray, shape (M,)
+        Indices into the original array of the kept representative.
+    voxel_map : dict
+        Mapping from voxel key (tuple of ints) -> representative index.
+    """
+    voxel_map = {}
+    eq_idx = []
+    for i, pt in enumerate(xyz):
+        key = tuple(np.round(pt / resolution).astype(int))
+        if key not in voxel_map:
+            eq_idx.append(i)
+            voxel_map[key] = len(eq_idx) - 1   # index into the EQUALIZED array
+    eq_idx = np.array(eq_idx, dtype=np.int64)
+    return xyz[eq_idx], eq_idx, voxel_map
+
+
+def compute_normals_and_curvature(xyz: np.ndarray, resolution: float = 0.1,
+                                  k: int = 30) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Estimate per-point normals and curvature using PCA on local neighborhoods.
+
+    The original LRGNet implementation uses a 3x3x3 voxel window search.
+    Here we offer both the exact voxel-window method and a k-NN fallback
+    (which is more robust for sparse clouds).
+
+    Parameters
+    ----------
+    xyz : np.ndarray, shape (N, 3)
+    resolution : float
+        Voxel size used to build the grid for fast neighbor lookup.
+    k : int
+        Minimum number of neighbors. If the voxel-window yields fewer,
+        we fall back to k-NN.
+
+    Returns
+    -------
+    normals : np.ndarray, shape (N, 3), float32
+        Absolute-value unit normals (always positive, matching original code).
+    curvature : np.ndarray, shape (N,), float32
+        Curvature = smallest_eigenvalue / sum_eigenvalues, in [0, 1].
+    """
+    n = len(xyz)
+    normals = np.zeros((n, 3), dtype=np.float32)
+    curvature = np.zeros(n, dtype=np.float32)
+
+    # Build voxel grid for fast lookups
+    voxel_grid = {}
+    for i, pt in enumerate(xyz):
+        key = tuple(np.round(pt / resolution).astype(int))
+        if key not in voxel_grid:
+            voxel_grid[key] = []
+        voxel_grid[key].append(i)
+
+    for i, pt in enumerate(xyz):
+        voxel = np.round(pt / resolution).astype(int)
+        # Search 3x3x3 voxel window (original method)
+        neighbors = []
+        for dx in (-1, 0, 1):
+            for dy in (-1, 0, 1):
+                for dz in (-1, 0, 1):
+                    key = (voxel[0]+dx, voxel[1]+dy, voxel[2]+dz)
+                    if key in voxel_grid:
+                        neighbors.extend(voxel_grid[key])
+        neighbors = np.array(neighbors, dtype=np.int64)
+
+        # Fallback to k-NN if too sparse
+        if len(neighbors) < k:
+            dists = np.linalg.norm(xyz - pt, axis=1)
+            neighbors = np.argsort(dists)[:k]
+
+        # Remove self
+        neighbors = neighbors[neighbors != i]
+        if len(neighbors) < 3:
+            normals[i] = [0, 0, 1]
+            curvature[i] = 0.0
+            continue
+
+        local_pts = xyz[neighbors] - pt
+        cov = local_pts.T @ local_pts / len(neighbors)
+        U, S, Vt = np.linalg.svd(cov)
+        # Normal = smallest eigenvector
+        normal = np.abs(Vt[2])
+        normals[i] = normal
+        curv = S[2] / (S[0] + S[1] + S[2] + 1e-8)
+        curvature[i] = curv
+
+    return normals, curvature
+
+
+def build_feature_vector(xyz: np.ndarray, rgb: Optional[np.ndarray],
+                         normals: np.ndarray, curvature: np.ndarray,
+                         room_bbox: Optional[np.ndarray] = None) -> np.ndarray:
+    """
+    Build the 13-channel feature vector used by LrgNet.
+
+    Feature layout (index: description):
+    0-2 : XYZ coordinates (absolute, in meters)
+    3-5 : Room-normalized coordinates (0..1 within scene bounding box)
+    6-8 : RGB colors, normalized to [-0.5, +0.5]
+    9-11: Normal vector (absolute value, always positive)
+    12  : Curvature (0..1)
+
+    Parameters
+    ----------
+    xyz : np.ndarray, shape (N, 3)
+    rgb : np.ndarray, shape (N, 3), uint8 or None
+    normals : np.ndarray, shape (N, 3)
+    curvature : np.ndarray, shape (N,)
+    room_bbox : np.ndarray, shape (2, 3), optional
+        Min/max corner of the room bounding box. Computed from xyz if omitted.
+
+    Returns
+    -------
+    features : np.ndarray, shape (N, 13), float32
+    """
+    n = len(xyz)
+    features = np.zeros((n, 13), dtype=np.float32)
+
+    # Channels 0-2: raw XYZ
+    features[:, :3] = xyz
+
+    # Channels 3-5: room-normalized coordinates
+    if room_bbox is None:
+        mins = xyz.min(axis=0)
+        maxs = xyz.max(axis=0)
+    else:
+        mins, maxs = room_bbox[0], room_bbox[1]
+    span = maxs - mins
+    span[span == 0] = 1.0
+    features[:, 3:6] = (xyz - mins) / span
+
+    # Channels 6-8: RGB normalized [-0.5, 0.5]
+    if rgb is not None:
+        rgb_f = rgb.astype(np.float32)
+        features[:, 6:9] = rgb_f / 255.0 - 0.5
+    else:
+        features[:, 6:9] = 0.0
+
+    # Channels 9-11: normals (already absolute)
+    features[:, 9:12] = normals
+
+    # Channel 12: curvature
+    features[:, 12] = curvature
+
+    return features