Upload learn_region_grow/utils.py

learn_region_grow/utils.py

@@ -0,0 +1,120 @@
+"""Utility helpers: sampling, centering, metrics."""
+
+import numpy as np
+from typing import Tuple
+
+
+def sample_or_pad(points: np.ndarray, target: int, seed: int = None) -> np.ndarray:
+    """
+    Sample exactly `target` points from `points` (N, C) with replacement if needed.
+
+    Parameters
+    ----------
+    points : np.ndarray, shape (N, C)
+    target : int
+        Number of points to return.
+    seed : int, optional
+        RNG seed for reproducibility.
+
+    Returns
+    -------
+    out : np.ndarray, shape (target, C)
+    """
+    n = len(points)
+    if n == 0:
+        return np.zeros((target, points.shape[1]), dtype=np.float32)
+    rng = np.random.default_rng(seed)
+    if n >= target:
+        idx = rng.choice(n, target, replace=False)
+    else:
+        idx = rng.choice(n, target, replace=True)
+    return points[idx]
+
+
+def center_features(inliers: np.ndarray, neighbors: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Translation / feature centering used by LrgNet.
+
+    Subtracts the median XY of the inlier set from XYZ of both inliers and neighbors.
+    Subtracts the median of channels 6+ (colors, normals, curvature) from both sets.
+
+    Parameters
+    ----------
+    inliers  : np.ndarray, shape (Ni, 13)
+    neighbors: np.ndarray, shape (Nn, 13)
+
+    Returns
+    -------
+    inliers_c, neighbors_c : same shapes, centered.
+    """
+    inliers = inliers.copy()
+    neighbors = neighbors.copy()
+
+    # Median XY of inliers
+    med_xy = np.median(inliers[:, :2], axis=0)
+    inliers[:, :2] -= med_xy
+    neighbors[:, :2] -= med_xy
+
+    # Median of channels 6+ (colors, normals, curvature) from inliers
+    med_feat = np.median(inliers[:, 6:], axis=0)
+    inliers[:, 6:] -= med_feat
+    neighbors[:, 6:] -= med_feat
+
+    return inliers, neighbors
+
+
+def cluster_metrics(pred_labels: np.ndarray, gt_labels: np.ndarray) -> dict:
+    """
+    Compute clustering quality metrics (NMI, AMI, ARS, Precision, Recall, IoU).
+
+    Parameters
+    ----------
+    pred_labels : np.ndarray, shape (N,)
+        Predicted instance IDs (0 = unlabeled / background).
+    gt_labels   : np.ndarray, shape (N,)
+        Ground-truth instance IDs.
+
+    Returns
+    -------
+    dict with keys: nmi, ami, ars, precision, recall, iou.
+    """
+    from sklearn.metrics import normalized_mutual_info_score, adjusted_mutual_info_score, adjusted_rand_score
+
+    mask = gt_labels >= 0
+    pred = pred_labels[mask]
+    gt = gt_labels[mask]
+
+    if len(pred) == 0:
+        return {"nmi": 0.0, "ami": 0.0, "ars": 0.0, "precision": 0.0, "recall": 0.0, "iou": 0.0}
+
+    nmi = normalized_mutual_info_score(gt, pred)
+    ami = adjusted_mutual_info_score(gt, pred)
+    ars = adjusted_rand_score(gt, pred)
+
+    # Precision / Recall / IoU over pairs
+    pred_pairs = set()
+    gt_pairs = set()
+    n = len(pred)
+    for i in range(n):
+        for j in range(i + 1, n):
+            if pred[i] == pred[j] and pred[i] >= 0:
+                pred_pairs.add((i, j))
+            if gt[i] == gt[j] and gt[i] >= 0:
+                gt_pairs.add((i, j))
+
+    tp = len(pred_pairs & gt_pairs)
+    fp = len(pred_pairs - gt_pairs)
+    fn = len(gt_pairs - pred_pairs)
+
+    precision = tp / (tp + fp + 1e-8)
+    recall = tp / (tp + fn + 1e-8)
+    iou = tp / (tp + fp + fn + 1e-8)
+
+    return {
+        "nmi": nmi,
+        "ami": ami,
+        "ars": ars,
+        "precision": precision,
+        "recall": recall,
+        "iou": iou,
+    }