Upload examples/segment_ply.py

examples/segment_ply.py

@@ -0,0 +1,85 @@
+"""Example: Segment a PLY point cloud with Point-SAM.
+
+Usage:
+    python segment_ply.py scene.ply 0.5 0.1 -0.2 --checkpoint model.safetensors
+"""
+
+import argparse
+
+import numpy as np
+from point_sam import PointSAM, load_pointcloud
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Segment a point cloud with Point-SAM")
+    parser.add_argument("ply_file", type=str, help="Path to .ply or .pcd file")
+    parser.add_argument("px", type=float, help="Prompt point x")
+    parser.add_argument("py", type=float, help="Prompt point y")
+    parser.add_argument("pz", type=float, help="Prompt point z")
+    parser.add_argument("--checkpoint", type=str, required=True, help="Path to .safetensors checkpoint")
+    parser.add_argument("--variant", type=str, default="large", choices=["large", "giant"])
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--output", type=str, default="mask.npy", help="Where to save the boolean mask")
+    args = parser.parse_args()
+
+    # 1. Load point cloud
+    coords, rgb, original_coords = load_pointcloud(args.ply_file, normalize=True)
+    print(f"Loaded {len(coords)} points from {args.ply_file}")
+
+    # 2. Load model
+    model = PointSAM.from_pretrained(
+        checkpoint_path=args.checkpoint,
+        variant=args.variant,
+        device=args.device,
+    )
+
+    # 3. Cache the point cloud for fast repeated queries
+    model.set_pointcloud(coords, rgb)
+
+    # 4. Run segmentation
+    prompt = [args.px, args.py, args.pz]
+    masks, iou_scores = model.predict(
+        coords=None,  # use cached
+        rgb=None,     # use cached
+        prompt_point=prompt,
+        prompt_label=1,  # foreground
+        multimask_output=True,
+    )
+
+    # 5. Pick best mask by IoU score
+    best_idx = int(np.argmax(iou_scores))
+    best_mask = masks[best_idx]
+    print(f"IoU scores: {iou_scores}")
+    print(f"Selected mask {best_idx} with {best_mask.sum()} points")
+
+    # 6. Save
+    np.save(args.output, best_mask)
+    print(f"Saved mask to {args.output}")
+
+    # 7. Optional: save segmented point cloud as PLY
+    segmented = original_coords[best_mask]
+    segmented_rgb = rgb[best_mask]
+    out_ply = args.output.replace(".npy", "_segmented.ply")
+    save_ascii_ply(out_ply, segmented, segmented_rgb)
+    print(f"Saved segmented point cloud to {out_ply}")
+
+
+def save_ascii_ply(path, coords, rgb):
+    """Save a point cloud as an ASCII PLY file."""
+    with open(path, "w") as f:
+        f.write("ply\n")
+        f.write("format ascii 1.0\n")
+        f.write(f"element vertex {len(coords)}\n")
+        f.write("property float x\n")
+        f.write("property float y\n")
+        f.write("property float z\n")
+        f.write("property uchar red\n")
+        f.write("property uchar green\n")
+        f.write("property uchar blue\n")
+        f.write("end_header\n")
+        for (x, y, z), (r, g, b) in zip(coords, rgb):
+            f.write(f"{x} {y} {z} {int(r)} {int(g)} {int(b)}\n")
+
+
+if __name__ == "__main__":
+    main()