align_clouds.py

#!/usr/bin/env python3
"""Align two point clouds with a similarity transform (rotation, translation,
uniform scale) and save the source cloud registered into the target's frame.

Written to accompany the OVER research note "How accurate is our automatic
360-camera reconstruction pipeline?" — running it on the published data
(fused.ply = photogrammetric reconstruction, all_color.pcd = LIDAR survey)
reproduces the note's headline numbers: scale factor ~0.9717 and a median
surface agreement of ~3.9 cm.

Method
------
1. Robust crop of far outliers (1st-99th percentile + 10% margin, registration
   only — the saved cloud is the untouched full-resolution source).
2. Global registration, RANSAC over FPFH features at 0.5 m voxels (skipped if
   --init is given).
3. Rigid ICP, coarse to medium thresholds (2.0 -> 0.5 m).
4. Scaled ICP (Umeyama point-to-point with free scale) at decreasing
   correspondence thresholds 0.50 / 0.25 / 0.12 / 0.06 m on 5 cm-voxel clouds.
   Freeing scale only at tight thresholds matters: loose correspondences bias
   the scale estimate low (the "shrink onto target" degeneracy).
5. The full-resolution source is transformed and saved as
   <source>_aligned_to_<target>.ply (colors preserved), together with a JSON
   holding the 4x4 similarity transform, its scale, and residual statistics.

Usage
-----
  python3 align_clouds.py SOURCE TARGET [--out-dir DIR] [--init transform.json]

The aligned PLY loads directly in CloudCompare / MeshLab next to the target.
"""
import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np
import open3d as o3d

reg = o3d.pipelines.registration
EST_RIGID = reg.TransformationEstimationPointToPoint(False)
EST_SCALED = reg.TransformationEstimationPointToPoint(True)


def log(*a):
    print("[align]", *a, flush=True)


def scale_of(T):
    return float(np.cbrt(abs(np.linalg.det(np.asarray(T)[:3, :3]))))


def crop_robust(pcd):
    """Drop far outliers so they cannot pollute voxel sizes / correspondences."""
    pts = np.asarray(pcd.points)
    lo, hi = np.percentile(pts, [1, 99], axis=0)
    margin = 0.10 * (hi - lo)
    m = np.all((pts >= lo - margin) & (pts <= hi + margin), axis=1)
    return pcd.select_by_index(np.flatnonzero(m))


def global_registration(src, tgt, voxel=0.5):
    """RANSAC over FPFH features; returns a rigid 4x4 or None."""
    def feat(p):
        d = p.voxel_down_sample(voxel)
        d.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=voxel * 2.5, max_nn=30))
        f = reg.compute_fpfh_feature(d, o3d.geometry.KDTreeSearchParamHybrid(radius=voxel * 5, max_nn=100))
        return d, f
    sd, sf = feat(src)
    td, tf = feat(tgt)
    log(f"global registration: {len(sd.points)} vs {len(td.points)} pts @ {voxel} m voxel")
    r = reg.registration_ransac_based_on_feature_matching(
        sd, td, sf, tf, True, voxel * 1.5, EST_RIGID, 3,
        [reg.CorrespondenceCheckerBasedOnEdgeLength(0.9),
         reg.CorrespondenceCheckerBasedOnDistance(voxel * 1.5)],
        reg.RANSACConvergenceCriteria(4_000_000, 0.9999))
    if len(r.correspondence_set) < 40:
        return None
    log(f"  RANSAC fitness {r.fitness:.3f}, {len(r.correspondence_set)} correspondences")
    return np.asarray(r.transformation)


def icp_cascade(src5, tgt5, T):
    """Rigid coarse-to-medium, then scaled fine ICP. src5/tgt5 are 5 cm-voxel clouds."""
    stages = [(2.00, EST_RIGID), (1.00, EST_RIGID), (0.50, EST_RIGID),
              (0.50, EST_SCALED), (0.25, EST_SCALED), (0.12, EST_SCALED), (0.06, EST_SCALED)]
    for thr, est in stages:
        r = reg.registration_icp(src5, tgt5, thr, T, est,
                                 reg.ICPConvergenceCriteria(max_iteration=80))
        T = np.asarray(r.transformation)
        kind = "scaled" if est is EST_SCALED else "rigid"
        log(f"  ICP {kind:6s} thr={thr:4.2f} m  fitness={r.fitness:.3f} "
            f"rmse={r.inlier_rmse:.4f}  scale={scale_of(T):.5f}")
    return T


def residual_stats(src5, tgt5, T):
    moved = o3d.geometry.PointCloud(src5)
    moved.transform(T)
    d = np.asarray(moved.compute_point_cloud_distance(tgt5))
    return {
        "n_points": int(len(d)),
        "median_m": float(np.median(d)),
        "p68_m": float(np.quantile(d, 0.68)),
        "p90_m": float(np.quantile(d, 0.90)),
        "within_5cm": float((d < 0.05).mean()),
        "within_10cm": float((d < 0.10).mean()),
    }


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("source", help="cloud to be moved (e.g. photogrammetric fused.ply)")
    ap.add_argument("target", help="reference cloud (e.g. LIDAR all_color.pcd)")
    ap.add_argument("--out-dir", default=".", help="where outputs are written")
    ap.add_argument("--init", help="JSON with a 'transform' 4x4 to skip global registration")
    ap.add_argument("--voxel", type=float, default=0.05,
                    help="registration/statistics voxel size in metres (default 0.05)")
    args = ap.parse_args()

    t0 = time.time()
    src_path, tgt_path = Path(args.source), Path(args.target)
    log("loading", src_path.name)
    src_full = o3d.io.read_point_cloud(str(src_path))
    log(f"  {len(src_full.points):,} points")
    log("loading", tgt_path.name)
    tgt_full = o3d.io.read_point_cloud(str(tgt_path))
    log(f"  {len(tgt_full.points):,} points")
    if len(src_full.points) < 1000 or len(tgt_full.points) < 1000:
        sys.exit("clouds too small")

    src = crop_robust(src_full)
    tgt = crop_robust(tgt_full)
    src5 = src.voxel_down_sample(args.voxel)
    tgt5 = tgt.voxel_down_sample(args.voxel)
    log(f"registration clouds @ {args.voxel} m voxel: "
        f"{len(src5.points):,} vs {len(tgt5.points):,} points")

    if args.init:
        T = np.array(json.loads(Path(args.init).read_text())["transform"], dtype=float)
        log("using provided init, scale", f"{scale_of(T):.5f}")
    else:
        T = global_registration(src, tgt)
        if T is None:
            sys.exit("global registration failed — provide a rough --init transform")

    T = icp_cascade(src5, tgt5, T)
    s = scale_of(T)
    if not 0.5 < s < 2.0:
        sys.exit(f"alignment degenerated (scale={s:.3f}) — check the inputs or provide --init")

    # residuals follow the note's protocol: FULL source (no outlier crop) at
    # the statistics voxel, distance to the nearest target point
    stats = residual_stats(src_full.voxel_down_sample(args.voxel), tgt5, T)
    log(f"FINAL scale {s:.6f} (inverse {1/s:.6f})")
    log(f"residuals: median {100*stats['median_m']:.1f} cm | p90 {100*stats['p90_m']:.1f} cm | "
        f"<5cm {100*stats['within_5cm']:.1f}% | <10cm {100*stats['within_10cm']:.1f}%")

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    out_ply = out_dir / f"{src_path.stem}_aligned_to_{tgt_path.stem}.ply"
    src_full.transform(T)
    o3d.io.write_point_cloud(str(out_ply), src_full)
    out_json = out_dir / f"{src_path.stem}_aligned_to_{tgt_path.stem}.transform.json"
    out_json.write_text(json.dumps({
        "source": src_path.name, "target": tgt_path.name,
        "transform": np.asarray(T).tolist(), "scale": s,
        "residuals": stats,
    }, indent=2))
    log("wrote", out_ply)
    log("wrote", out_json)
    log(f"done in {time.time()-t0:.0f} s — load {out_ply.name} together with "
        f"{tgt_path.name} in CloudCompare to inspect the overlay")


if __name__ == "__main__":
    main()