Upload scripts/segment.py

scripts/segment.py

@@ -0,0 +1,96 @@
+"""
+segment.py
+==========
+
+Command-line script to run learned region growing on a PLY / PCD point cloud.
+
+Example:
+    python segment.py --input scene.ply --ckpt checkpoints/best_model.pt \
+                      --output segmented_scene.ply --device cuda
+"""
+
+import argparse
+import torch
+import numpy as np
+from pathlib import Path
+
+from learn_region_grow.io import load_point_cloud, save_ply
+from learn_region_grow.preprocess import voxel_equalize, compute_normals_and_curvature, build_feature_vector
+from learn_region_grow.lrg_net import LrgNet
+from learn_region_grow.growing import RegionGrower
+
+
+def main():
+    parser = argparse.ArgumentParser(description="LRGNet: Learned Region Growing on a Point Cloud")
+    parser.add_argument("--input", required=True, help="Input .ply or .pcd file")
+    parser.add_argument("--ckpt", required=True, help="PyTorch checkpoint (.pt)")
+    parser.add_argument("--output", default="output.ply", help="Output segmented PLY")
+    parser.add_argument("--device", default="cuda", help="cuda or cpu")
+    parser.add_argument("--resolution", type=float, default=0.1, help="Voxel resolution in meters")
+    parser.add_argument("--lite", type=int, default=0, choices=[0,1,2], help="Lite model variant")
+    parser.add_argument("--stochastic", action="store_true", default=True, help="Use stochastic growing (default)")
+    parser.add_argument("--deterministic", dest="stochastic", action="store_false", help="Use deterministic thresholding")
+    parser.add_argument("--add_threshold", type=float, default=0.5)
+    parser.add_argument("--remove_threshold", type=float, default=0.5)
+    parser.add_argument("--cluster_threshold", type=int, default=10)
+    args = parser.parse_args()
+
+    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
+    print(f"Loading point cloud from {args.input} ...")
+    xyz, rgb, normals_input = load_point_cloud(args.input)
+    print(f"  {len(xyz)} points loaded")
+
+    # Voxel equalization
+    print(f"Voxel equalization (resolution={args.resolution}m) ...")
+    eq_xyz, eq_idx, voxel_map = voxel_equalize(xyz, args.resolution)
+    eq_rgb = rgb[eq_idx] if rgb is not None else None
+    print(f"  {len(eq_xyz)} points after equalization")
+
+    # Normals / curvature
+    if normals_input is not None:
+        normals = np.abs(normals_input[eq_idx])
+        curvature = np.zeros(len(eq_xyz), dtype=np.float32)  # could estimate from normals
+        # Better: still compute curvature from PCA
+        _, curvature = compute_normals_and_curvature(eq_xyz, args.resolution)
+    else:
+        print("Computing normals & curvature via PCA ...")
+        normals, curvature = compute_normals_and_curvature(eq_xyz, args.resolution)
+
+    # Build feature vector
+    features = build_feature_vector(eq_xyz, eq_rgb, normals, curvature)
+    print(f"Feature vector shape: {features.shape}")
+
+    # Load model
+    print(f"Loading checkpoint {args.ckpt} ...")
+    model = LrgNet(in_channels=13, lite=args.lite)
+    model.load_state_dict(torch.load(args.ckpt, map_location=device))
+    model.to(device)
+    print("Model loaded.")
+
+    # Run region growing
+    grower = RegionGrower(
+        model, device,
+        add_threshold=args.add_threshold,
+        remove_threshold=args.remove_threshold,
+        cluster_threshold=args.cluster_threshold,
+        stochastic=args.stochastic,
+    )
+    print("Running learned region growing ...")
+    labels = grower.grow(eq_xyz, features, voxel_map, args.resolution)
+    n_instances = len(np.unique(labels[labels >= 0]))
+    print(f"Segmented into {n_instances} instances")
+
+    # Expand labels back to original dense cloud via nearest equalized neighbor
+    from scipy.spatial import cKDTree
+    tree = cKDTree(eq_xyz)
+    _, nn = tree.query(xyz)
+    full_labels = labels[nn]
+
+    # Save
+    print(f"Saving output to {args.output} ...")
+    save_ply(args.output, xyz, labels=full_labels)
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()