Upload scripts/stage_data.py

scripts/stage_data.py

@@ -0,0 +1,61 @@
+"""
+stage_data.py
+=============
+
+Generate staged training data from labeled point clouds.
+
+Requires point clouds with instance labels. For PLY/PCD files without labels,
+you can provide a separate label file (e.g. .npy with shape (N,)).
+
+Example:
+    python stage_data.py --input scene.ply --labels scene_labels.npy \
+                         --output scene_staged.h5 --resolution 0.1
+"""
+
+import argparse
+import numpy as np
+from pathlib import Path
+
+from learn_region_grow.io import load_point_cloud
+from learn_region_grow.stage_data import stage_labeled_cloud, save_staged_h5
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Stage labeled point clouds into training H5")
+    parser.add_argument("--input", required=True, help="Input .ply or .pcd")
+    parser.add_argument("--labels", default=None, help="Optional .npy file with (N,) instance labels")
+    parser.add_argument("--output", required=True, help="Output .h5 file")
+    parser.add_argument("--resolution", type=float, default=0.1)
+    parser.add_argument("--add_mistake_prob", type=float, default=0.2)
+    parser.add_argument("--remove_mistake_prob", type=float, default=0.2)
+    parser.add_argument("--seeds_per_instance", type=int, default=5)
+    parser.add_argument("--max_steps_per_seed", type=int, default=20)
+    args = parser.parse_args()
+
+    xyz, rgb, _ = load_point_cloud(args.input)
+    print(f"Loaded {len(xyz)} points")
+
+    if args.labels:
+        labels = np.load(args.labels)
+    else:
+        # Try to read labels from PLY / PCD if present
+        # For now, raise
+        raise ValueError("--labels is required. Provide a .npy file with shape (N,) integer instance IDs.")
+
+    print(f"Staging data ...")
+    inliers, neighbors, add_lbl, rmv_lbl = stage_labeled_cloud(
+        xyz, rgb, labels,
+        add_mistake_prob=args.add_mistake_prob,
+        remove_mistake_prob=args.remove_mistake_prob,
+        resolution=args.resolution,
+        seeds_per_instance=args.seeds_per_instance,
+        max_steps_per_seed=args.max_steps_per_seed,
+    )
+
+    print(f"Generated {len(inliers)} training tuples")
+    save_staged_h5(args.output, inliers, neighbors, add_lbl, rmv_lbl)
+    print(f"Saved to {args.output}")
+
+
+if __name__ == "__main__":
+    main()