scripts/find_placement_surface.py

#!/usr/bin/env python
"""Fit the top surface of a masked object (the book stack) as a robot placement target.

This is deliberately its own script rather than a sixth `run_object_pipeline.py`
stage: none of stages 2-5 (mesh / align / act / render) apply here -- there is
no need to fit a mesh or animate a trajectory for a static stack of books we
only want a flat place-target on top of. All that is needed is stage 1's mask
plus the depth-derived point cloud stage 3 already knows how to build
(`fpgm.objects.align.object_point_cloud`), and a plane fit on top of it.

Method:
    1. `object_point_cloud` turns the book-stack mask + frame-0 depth into a
       world-frame (robot-base, `panda_link0`) point cloud, with the same
       robust MAD depth-outlier rejection stage 3 relies on -- a handful of
       leaked background pixels would otherwise drag the fit metres away (see
       that function's own docstring for the measured 0.39 m example).
    2. The book stack is not flat -- it is several books of different
       thickness with a near-vertical front/page edge below the top cover, so
       a plane fit over the *whole* cloud averages the stair-stepped sides
       into a tilted mess. Only a *thin* band at the very top of the cloud's
       world-Z range is the top book's actual flat cover; empirically on this
       scene that band is only ~8% of the full stack height (see
       `--top-band-frac`'s docstring) -- looser than that and the band starts
       including the front edge, which has more points than the true top face
       at this camera's grazing angle and pulls a plain fit off-horizontal
       (this was observed directly while building this script; see the
       validation-stance note below).
    3. RANSAC first, gated to near-horizontal candidates only (see
       `--min-normal-up`) since a placement target must be a surface an
       object can rest on under gravity, not merely whatever plane the data
       supports most -- then one SVD least-squares refit on the RANSAC inlier
       set, which is the actual best-fit plane through the points RANSAC
       identified as "the top face" rather than through an arbitrary 3-point
       sample.

Validation stance (read before trusting the output):
    A prior mesh-fitting bug in this project reported rmse=1.96cm,
    inlier_fraction=99% for a mesh that was 14x too large in volume, because
    point-to-surface RMSE is a ONE-SIDED metric -- an oversized surface that
    encloses the observed points scores just as well as a correctly-sized one.
    The same failure mode applies here: a plane can have a tiny RMS residual
    while its RANSAC inlier set is a tiny, unrepresentative patch of the real
    top face. So this script always reports `surface_extent_m` (the fitted
    plane's own inlier footprint) *and* `cloud_extent_m` /
    `cloud_footprint_uv_m` (the full observed cloud's extent, the second one
    projected through the identical plane basis so the comparison is
    apples-to-apples) side by side, and logs a loud warning if the plane's
    footprint covers under ~30% of what the full cloud's footprint suggests is
    there. A human should also look at the debug overlay PNG -- these numbers
    alone cannot distinguish "SAM correctly masked a small book cover" from
    "the plane fit collapsed onto a corner".

Usage:
    PYTHONPATH=src python scripts/find_placement_surface.py \\
        --episode <uuid> --camera ext1 --clip 0:11 \\
        --mask-outputs outputs/<uuid>/objects_books/0_11 \\
        --out outputs/<uuid>/objects_books/0_11/placement
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import sys
from pathlib import Path

import cv2
import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

from fpgm.objects.align import object_point_cloud  # noqa: E402
from fpgm.objects.scene import draw_point_cloud, project_point_cloud  # noqa: E402
from fpgm.utils.io import ensure_dir  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

# Reuse run_object_pipeline.py's private clip-context/metadata/depth loaders
# instead of duplicating them -- same trick as scripts/mask_object_manual.py
# and scripts/measure_push_tracked.py. The sys.modules registration before
# exec_module is required: run_object_pipeline.py's ClipContext is a
# @dataclass, and dataclass decorators look the defining module up in
# sys.modules while they run, which fails if the module isn't registered yet.
_spec = importlib.util.spec_from_file_location(
    "run_object_pipeline", REPO_ROOT / "scripts" / "run_object_pipeline.py"
)
rop = importlib.util.module_from_spec(_spec)
sys.modules[_spec.name] = rop
_spec.loader.exec_module(rop)

logger = get_logger("find_placement_surface")

#: Robot-base (panda_link0) convention used throughout this pipeline: +Z is up.
#: See fpgm/objects/scene.py's module-level constant for the same assumption.
_WORLD_UP = np.array([0.0, 0.0, 1.0])

#: Below this, a plane footprint is flagged as suspiciously small relative to
#: the full observed cloud's footprint through the same basis -- see the
#: module docstring's validation-stance note.
_FOOTPRINT_RATIO_WARN = 0.30


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--episode", required=True)
    p.add_argument("--camera", default="ext1", choices=["ext1", "ext2"])
    p.add_argument("--clip", default="0:11")
    p.add_argument("--mask-outputs", type=Path, required=True,
                    help="the --outputs directory previously passed to mask_object_manual.py")
    p.add_argument("--top-band-frac", type=float, default=0.08,
                    help="fraction of the cloud's world-Z range, from the top, treated as "
                    "'the top surface' before plane fitting (see module docstring). Tuned "
                    "empirically on the book-stack scene this script was built for: a stack's "
                    "actual flat top cover is a thin sliver at the very top (~5-10%% of the full "
                    "stack height here); anything looser (measured up to ~0.30 on this scene) "
                    "starts pulling in the topmost book's near-vertical front/page edge, which "
                    "has *more* points than the true top face at this camera's grazing angle and "
                    "so drags a plain RANSAC fit off-horizontal. Re-tune per scene with "
                    "--min-normal-up as the tell: if the reported normal isn't close to vertical, "
                    "shrink this until it is (see that flag's own docstring)")
    p.add_argument("--ransac-iters", type=int, default=2000)
    p.add_argument("--ransac-thresh-m", type=float, default=0.006,
                    help="inlier distance-to-plane threshold, metres; ~ the depth noise floor "
                    "for a workbench-distance RGBD capture, not a tuned-for-this-object value")
    p.add_argument("--min-normal-up", type=float, default=0.7,
                    help="minimum |normal . world_up| for a candidate plane to be considered at "
                    "all -- a placement TARGET must be a surface an object can rest on under "
                    "gravity (near-horizontal), and without this a plain max-inliers RANSAC can "
                    "(and, on this scene, does) lock onto a larger near-vertical face -- e.g. a "
                    "book's front cover/page-edge -- instead of its horizontal top, because that "
                    "face happens to carry more points at this camera's grazing viewing angle on "
                    "the true top. 0.7 ~= within 45 degrees of vertical")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--out", type=Path, default=None,
                    help="output directory (default: <mask-outputs>/objects/2_placement)")
    return p.parse_args()


def _svd_plane(points: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Least-squares plane through `points`: (normal, centroid, in-plane basis).

    The centroid of the input points always lies exactly on the returned
    plane -- it is the point the plane is fit *through* -- so callers that want
    "a point on the plane" never need a separate projection step for the
    centroid specifically.

    `basis` is `(2, 3)`: the two largest right-singular vectors of the
    centered points, i.e. the in-plane directions of greatest spread. Used
    both as the projection axes for footprint/extent measurements and to keep
    that footprint's orientation meaningful (aligned with the surface's own
    long/short axes) rather than arbitrary world XY.
    """
    centroid = points.mean(axis=0)
    centered = points - centroid
    _, _, vt = np.linalg.svd(centered, full_matrices=False)
    normal = vt[-1]
    basis = vt[:2]
    return normal, centroid, basis


def _fit_top_plane(
    band_points: np.ndarray,
    thresh_m: float,
    n_iters: int,
    rng: np.random.Generator,
    min_normal_up: float,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, float]:
    """RANSAC-then-refit plane fit. Returns (normal, centroid, basis, inlier_mask, rms_m).

    RANSAC only ever needs to decide *which points are inliers* -- the 3-point
    candidate plane's own accuracy does not matter, since the final normal and
    centroid always come from an SVD refit over whichever inlier set won (see
    below). This matters because a 3-point plane is exactly as noise-sensitive
    as the single depth pixels it is built from; the refit is what actually
    uses the redundancy in the ~hundreds of top-band points to average that
    noise out.

    `min_normal_up` rejects candidate 3-point planes that are not near-horizontal
    *before* they are scored by inlier count. Plain max-inlier RANSAC picks
    whichever plane the data supports most, with no notion of "supports an
    object resting on it" -- on this script's own first (unconstrained) run
    against the book-stack scene, that picked a near-vertical face (front
    cover/page-edge) over the true horizontal top cover, because the vertical
    face happened to carry more consistent points at this camera's grazing
    viewing angle on the actual top. The gate encodes the one prior that
    actually matters for a placement target -- gravity -- directly, instead of
    hoping "top band by Z" alone disambiguates it.
    """
    n = band_points.shape[0]
    if n < 3:
        raise ValueError(f"_fit_top_plane: need >=3 points in the top band, got {n}")

    best_inliers: np.ndarray | None = None
    best_count = -1
    n_tried = 0
    for _ in range(n_iters):
        idx = rng.choice(n, size=3, replace=False)
        p0, p1, p2 = band_points[idx]
        normal = np.cross(p1 - p0, p2 - p0)
        norm = np.linalg.norm(normal)
        if norm < 1e-12:
            continue  # degenerate (near-collinear) sample; skip rather than divide by ~0
        normal = normal / norm
        if abs(float(np.dot(normal, _WORLD_UP))) < min_normal_up:
            continue  # not near-horizontal; not a candidate resting surface, see docstring
        n_tried += 1
        d = -normal @ p0
        dist = band_points @ normal + d
        inliers = np.abs(dist) <= thresh_m
        count = int(inliers.sum())
        if count > best_count:
            best_count, best_inliers = count, inliers

    if best_inliers is None or best_count < 3:
        raise ValueError(
            f"_fit_top_plane: found no near-horizontal (|normal.up| >= {min_normal_up}) 3-point "
            f"plane candidate with >=3 inliers out of {n_iters} samples ({n_tried} passed the "
            "horizontality gate) -- the top band likely does not contain a real resting surface; "
            "try a different --top-band-frac or inspect the mask"
        )

    # One SVD refit through the winning inlier set, then one more inlier pass
    # against *that* refined plane (a handful of points near the threshold
    # boundary can flip in/out once the plane itself is no longer a noisy
    # 3-point guess) before the final refit that is actually reported.
    normal, centroid, basis = _svd_plane(band_points[best_inliers])
    dist = (band_points - centroid) @ normal
    inliers2 = np.abs(dist) <= thresh_m
    normal, centroid, basis = _svd_plane(band_points[inliers2])

    dist_final = (band_points - centroid) @ normal
    rms = float(np.sqrt(np.mean(dist_final[inliers2] ** 2)))
    if abs(float(np.dot(normal, _WORLD_UP))) < min_normal_up:
        raise ValueError(
            f"_fit_top_plane: refit through the winning inlier set drifted back to a non-horizontal "
            f"normal {normal.round(3).tolist()} (|.up|={abs(float(np.dot(normal, _WORLD_UP))):.3f} < "
            f"{min_normal_up}) -- the inlier set the horizontality gate accepted was not actually "
            "coherent; do not trust this fit"
        )
    return normal, centroid, basis, inliers2, rms


def _in_plane_extent(points: np.ndarray, centroid: np.ndarray, basis: np.ndarray) -> np.ndarray:
    """`(2,)` [u_extent, v_extent]: `points`' spread along the plane's own basis axes."""
    projected = (points - centroid) @ basis.T
    return projected.max(axis=0) - projected.min(axis=0)


def main() -> int:
    args = parse_args()
    setup_logging()

    mask_dir = rop.stage_dir(args.mask_outputs, "mask")
    meta = json.loads((mask_dir / "meta.json").read_text())
    with np.load(mask_dir / "mask_frame0.npz") as npz:
        mask_full = npz["mask"].astype(bool)
    video_w, video_h = meta["video_resolution"]
    logger.info("loaded mask: prompt=%r, %d px on, video_resolution=%dx%d",
                meta["prompt"], int(mask_full.sum()), video_w, video_h)

    out_dir = ensure_dir(args.out) if args.out is not None else ensure_dir(
        Path(args.mask_outputs) / "objects" / "2_placement"
    )

    ctx = rop._load_clip_context(args.episode, args.camera, args.clip)
    depth = rop._load_initial_depth(args.episode, args.clip, meta["camera_serial"])
    adapter = rop._DepthRgbAdapter(depth, ctx.clip.initial_rgb)

    points_world, colors = object_point_cloud(adapter, ctx.camera_annot, mask_full, (video_h, video_w))
    n_cloud = points_world.shape[0]
    cloud_centroid = points_world.mean(axis=0)
    cloud_extent = points_world.max(axis=0) - points_world.min(axis=0)
    logger.info(
        "object cloud: %d points, extent=%s m, centroid=%s m",
        n_cloud, np.round(cloud_extent, 4).tolist(), np.round(cloud_centroid, 4).tolist(),
    )

    z = points_world[:, 2]
    z_min, z_max = float(z.min()), float(z.max())
    z_threshold = z_max - args.top_band_frac * (z_max - z_min)
    band_mask = z >= z_threshold
    band_points = points_world[band_mask]
    logger.info(
        "top band: z in [%.4f, %.4f] m (top %.0f%% of full z-range [%.4f, %.4f]) -> %d/%d points",
        z_threshold, z_max, args.top_band_frac * 100, z_min, z_max, int(band_mask.sum()), n_cloud,
    )

    rng = np.random.default_rng(args.seed)
    normal, position, basis, inlier_mask_band, rms_m = _fit_top_plane(
        band_points, args.ransac_thresh_m, args.ransac_iters, rng, args.min_normal_up
    )
    n_inliers = int(inlier_mask_band.sum())
    inlier_points = band_points[inlier_mask_band]

    # PCA/SVD gives the normal's axis but not its sign; "up" in this pipeline's
    # world frame is a fixed, known convention (+Z, panda_link0), not something
    # to infer from the data, so it is enforced directly rather than left to
    # whichever of the two SVD sign conventions happened to come out.
    if np.dot(normal, _WORLD_UP) < 0.0:
        normal = -normal
        basis = np.array([basis[0], -basis[1]])  # keep basis right-handed with the flipped normal

    surface_extent = _in_plane_extent(inlier_points, position, basis)
    cloud_footprint_uv = _in_plane_extent(points_world, position, basis)

    ratios = surface_extent / np.maximum(cloud_footprint_uv, 1e-9)
    extent_flag = bool(np.any(ratios < _FOOTPRINT_RATIO_WARN))
    logger.info(
        "plane fit: n_inliers=%d/%d, rms=%.4f m, normal=%s, position=%s",
        n_inliers, band_points.shape[0], rms_m, np.round(normal, 4).tolist(), np.round(position, 4).tolist(),
    )
    logger.info(
        "extent check: surface_extent_m(in-plane)=%s vs cloud_footprint_uv_m(same basis)=%s "
        "(ratios=%s)",
        np.round(surface_extent, 4).tolist(), np.round(cloud_footprint_uv, 4).tolist(),
        np.round(ratios, 3).tolist(),
    )
    if extent_flag:
        logger.warning(
            "!! plane footprint covers only %.0f%% of the full cloud's footprint along at least "
            "one axis (threshold %.0f%%) -- this is exactly the one-sided-RMSE failure mode this "
            "script was told to guard against (see module docstring); inspect the debug overlay "
            "before trusting this placement target.",
            100.0 * float(ratios.min()), _FOOTPRINT_RATIO_WARN * 100.0,
        )

    result = {
        "episode": args.episode,
        "camera": args.camera,
        "clip": args.clip,
        "prompt": meta["prompt"],
        "position_m": position.tolist(),
        "normal": normal.tolist(),
        "plane_rms_m": rms_m,
        "n_inliers": n_inliers,
        "n_band_points": int(band_points.shape[0]),
        "n_cloud_points": n_cloud,
        "surface_extent_m": surface_extent.tolist(),
        "cloud_extent_m": cloud_extent.tolist(),
        "cloud_footprint_uv_m": cloud_footprint_uv.tolist(),
        "cloud_centroid_m": cloud_centroid.tolist(),
        "z_band_used": [z_threshold, z_max],
        "footprint_ratio_min": float(ratios.min()),
        "footprint_ratio_flagged": extent_flag,
        "ransac_thresh_m": args.ransac_thresh_m,
        "top_band_frac": args.top_band_frac,
    }
    out_path = out_dir / "placement_target.json"
    out_path.write_text(json.dumps(result, indent=2))
    logger.info("wrote %s", out_path)

    # --- debug overlay: project the full cloud (context), the plane inliers
    # (what the fit actually used), and the placement point back onto frame 0.
    frame0_bgr = cv2.imread(str(mask_dir / "frame0_bgr.png"))
    if frame0_bgr is None:
        raise FileNotFoundError(f"could not read {mask_dir / 'frame0_bgr.png'}")
    camera_video = ctx.camera_annot.rescaled(video_w, video_h)

    overlay = frame0_bgr.copy()
    uv_cloud, colors_cloud = project_point_cloud(points_world, colors, camera_video, (video_h, video_w))
    overlay = draw_point_cloud(overlay, uv_cloud, colors_cloud, radius=1)

    green = np.tile(np.array([[60, 220, 60]], dtype=np.uint8), (inlier_points.shape[0], 1))
    uv_inliers, colors_inliers = project_point_cloud(inlier_points, green, camera_video, (video_h, video_w))
    overlay = draw_point_cloud(overlay, uv_inliers, colors_inliers, radius=2)

    uv_pos, depth_pos = camera_video.project(position[None, :])
    if depth_pos[0] > 0:
        center = tuple(np.round(uv_pos[0]).astype(int))
        cv2.drawMarker(overlay, center, (0, 0, 255), markerType=cv2.MARKER_CROSS,
                        markerSize=24, thickness=2, line_type=cv2.LINE_AA)
        cv2.circle(overlay, center, 10, (0, 0, 255), 2, lineType=cv2.LINE_AA)
        # A short segment along the outward normal, so the overlay also shows
        # whether "up" came out pointing the right way, not just "somewhere".
        tip_world = position + 0.05 * normal
        uv_tip, depth_tip = camera_video.project(tip_world[None, :])
        if depth_tip[0] > 0:
            cv2.line(overlay, center, tuple(np.round(uv_tip[0]).astype(int)),
                      (0, 0, 255), 2, lineType=cv2.LINE_AA)

    overlay_path = out_dir / "placement_overlay.png"
    cv2.imwrite(str(overlay_path), overlay)
    logger.info("wrote %s", overlay_path)

    print(json.dumps(result, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())