src/fpgm/objects/align.py

"""Stage 3: mask x depth -> world point cloud -> scale + 6-DOF mesh pose.

Three things make this stage trickier than "unproject the mask and call it a
day":

* A segmentation mask that leaks even a handful of background pixels drags
  the centroid metres away from the object -- :func:`object_point_cloud`'s
  outlier rejection exists because of exactly this, observed directly: a leaky
  mask on a small LEGO brick produced a 0.39 m bounding box for a 0.03 m
  object.
* The depth camera only ever sees the near side of the object -- the point
  cloud is a single-view *shell*, not a full surface -- which biases a naive
  full-mesh ICP fit. See :func:`fit_mesh_to_points`'s docstring for how that is
  handled.
* For objects at centimetre scale, the depth map's *own noise* can exceed the
  object's size, so the point cloud's spatial extent is not a measurement of
  the object at all -- it is a measurement of depth noise. Measured directly
  on a 2x4 LEGO brick (true size ~0.032 x 0.016 x 0.019 m) at 0.63 m: depth
  MAD inside its (correct, tight) SAM mask was 0.0265 m, and the resulting
  point cloud spanned 0.215 x 0.185 x 0.121 m -- wrong by up to 7x, and no
  amount of outlier trimming recovers the true extent from it (+/-1cm
  trimming around the median still leaves 0.052 x 0.091 x 0.056 m). The mask
  itself, in contrast, was segmented perfectly (1664 px, a tight 82x81 px
  bbox); its *angular* extent at the (reliable) median depth gives the right
  answer. :func:`fit_mesh_to_points` therefore fits position, scale, and
  orientation from whichever measurement actually supports each one -- see
  its docstring.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Protocol

import cv2
import numpy as np
import trimesh
from scipy.spatial import cKDTree

from fpgm.geometry.camera import Camera
from fpgm.objects.types import ObjectAlignment, ObjectMesh, StageArtifacts
from fpgm.types import FpgmError
from fpgm.utils.io import ensure_dir
from fpgm.utils.logging import get_logger
from fpgm.viz.overlays import color_for

logger = get_logger(__name__)


class AlignmentError(FpgmError):
    """Raised when the observed cloud or requested fit is too degenerate to trust.

    Deliberately distinct from a bare :class:`ValueError`: catching this one
    error class is how a caller distinguishes "this object's geometry could not
    be aligned" from an ordinary programming mistake (wrong argument types etc).
    """


class DepthRgbSource(Protocol):
    """Minimal per-clip contract :func:`object_point_cloud` needs.

    :class:`fpgm.types.SceneFlowClip` carries per-*query-point* scene-flow
    annotations, not the dense per-pixel ``initial_depth``/``initial_rgb``
    arrays this function needs -- those live as sibling h5 datasets of
    ``intrinsic``/``extrinsic`` inside the same ``"camera_<serial>_ext"``
    group. Rather than widen the shared ``SceneFlowClip`` contract for one
    stage, this is a structural (duck-typed) protocol: any object exposing
    these two attributes at the camera's resolution works, e.g. a thin wrapper
    around the raw h5 group.
    """

    initial_depth: np.ndarray  # (H, W) uint16, millimetres
    initial_rgb: np.ndarray | None  # (H, W, 3) uint8, or None


_DEPTH_MM_TO_M = 1.0 / 1000.0
_MIN_FIT_POINTS = 6
#: A cloud is treated as near-collinear (orientation undefined, unrecoverable)
#: once its second-largest PCA eigenvalue is this small relative to its
#: largest -- the largest-alone can be huge for an elongated-but-well-posed
#: object, so the *ratio between the top two* is what actually signals "no
#: second direction". This is a hard failure (raises); see
#: _LOW_STRUCTURE_EIGVAL_RATIO for the softer "don't trust it" signal.
_COLLINEAR_EIGVAL_RATIO = 1e-6
#: A cloud with a near-zero *third* eigenvalue (relative to the first) has at
#: most 2 well-defined directions -- e.g. a near-planar shell -- so no 3-axis
#: orientation can be trusted, even though it is not collinear enough to
#: outright raise. Looser than _COLLINEAR_EIGVAL_RATIO on purpose: this is a
#: "fall back quietly" signal, not a "this call is nonsense" one.
_LOW_STRUCTURE_EIGVAL_RATIO = 1e-4
#: The four proper-rotation (det = +1) sign patterns for a right-handed axis
#: triad -- PCA's covariance is blind to axis handedness, so all four are
#: tried and scored against the cloud before ICP starts.
_SIGN_CANDIDATES = ((1, 1, 1), (1, -1, -1), (-1, 1, -1), (-1, -1, 1))


def object_point_cloud(
    clip: DepthRgbSource,
    camera: Camera,
    mask_full: np.ndarray,
    frame_shape: tuple[int, int],
    max_depth_span_m: float = 0.5,
    mad_k: float = 3.0,
) -> tuple[np.ndarray, np.ndarray]:
    """Unproject the masked, depth-valid pixels of one frame into world points.

    Args:
        clip: Exposes ``initial_depth`` (mm, uint16) and optionally
            ``initial_rgb`` at ``camera``'s resolution -- see
            :class:`DepthRgbSource`.
        camera: Camera already :meth:`~fpgm.geometry.camera.Camera.rescaled` to
            ``clip.initial_depth``'s resolution.
        mask_full: ``(H, W)`` boolean object mask, at ``frame_shape``
            resolution (typically the video's full resolution, coarser than
            the depth map).
        frame_shape: ``(H, W)`` the resolution ``mask_full`` is expected to be
            at -- a cheap sanity check that catches a caller passing a
            mismatched mask rather than silently misaligning it.
        max_depth_span_m: Hard cap on ``|depth - median_depth|`` for a pixel to
            be kept; see the outlier-rejection note below.
        mad_k: Robust z-score cutoff (multiples of scaled MAD) for outlier
            rejection.

    Returns:
        ``(points_world, colors)``: ``(N, 3)`` float64 world points and
        ``(N, 3)`` uint8 colours (zeros if ``clip.initial_rgb`` is ``None``).

    Raises:
        ValueError: If ``mask_full``'s shape disagrees with ``frame_shape``, or
            no pixels survive masking / outlier rejection.

    Outlier rejection:
        A segmentation mask that leaks even a handful of background pixels
        drags the centroid metres away from the object -- this was observed
        directly: a leaky mask on a ~0.03 m LEGO brick produced a 0.39 m
        bounding box, because a few background pixels at very different depth
        pulled the extent far past the true object. Depth is therefore kept
        only within a robust band (``median +/- mad_k`` scaled MAD) and never
        farther than ``max_depth_span_m / 2`` from the median, so a mask leak
        is rejected as an outlier instead of being silently baked into the
        returned cloud.
    """
    mask_full = np.asarray(mask_full, dtype=bool)
    if mask_full.shape != tuple(frame_shape):
        raise ValueError(
            f"object_point_cloud: mask_full shape {mask_full.shape} != "
            f"declared frame_shape {tuple(frame_shape)}"
        )

    depth_mm = np.asarray(clip.initial_depth)
    camera.sanity_check_resolution(depth_mm.shape)
    depth_h, depth_w = depth_mm.shape[:2]

    # cv2.resize with anything but nearest-neighbour would blend foreground and
    # background labels at the boundary into a fractional "mask value"; a bool
    # mask has no such thing, so nearest is the only interpolation that keeps
    # the result a mask.
    mask_small = cv2.resize(
        mask_full.astype(np.uint8), (depth_w, depth_h), interpolation=cv2.INTER_NEAREST
    ).astype(bool)

    depth_m = depth_mm.astype(np.float64) * _DEPTH_MM_TO_M
    valid = mask_small & (depth_m > 0)
    if not valid.any():
        raise ValueError("object_point_cloud: no valid (masked, depth>0) pixels")

    ys, xs = np.nonzero(valid)
    depths = depth_m[ys, xs]

    keep = _robust_depth_mask(depths, max_depth_span_m, mad_k)
    n_rejected = int((~keep).sum())
    if n_rejected:
        logger.info(
            "object_point_cloud: rejected %d/%d masked pixels as depth outliers "
            "(median=%.3fm)",
            n_rejected,
            depths.shape[0],
            float(np.median(depths)),
        )
    ys, xs, depths = ys[keep], xs[keep], depths[keep]
    if depths.size == 0:
        raise ValueError("object_point_cloud: all masked pixels rejected as depth outliers")

    uv = np.stack([xs, ys], axis=1).astype(np.float64)
    points_world = camera.unproject(uv, depths)

    if clip.initial_rgb is not None:
        colors = np.asarray(clip.initial_rgb)[ys, xs].astype(np.uint8)
    else:
        colors = np.zeros((depths.size, 3), dtype=np.uint8)

    return points_world, colors


def _robust_depth_mask(depths: np.ndarray, max_depth_span_m: float, mad_k: float) -> np.ndarray:
    median = np.median(depths)
    abs_dev = np.abs(depths - median)
    mad = np.median(abs_dev)
    half_span = max_depth_span_m / 2.0
    if mad < 1e-9:
        # Degenerate/near-constant depth: MAD carries no information (would
        # reject everything but the exact median), fall back to the absolute
        # span cap alone.
        return abs_dev <= half_span
    robust_std = 1.4826 * mad  # scales MAD to a normal-equivalent std
    return (abs_dev <= mad_k * robust_std) & (abs_dev <= half_span)


def _resize_mask_nearest(mask: np.ndarray, height: int, width: int) -> np.ndarray:
    """Nearest-neighbour resize a bool mask to ``(height, width)`` (no-op if already that shape).

    Nearest is the only interpolation that keeps a bool mask a mask -- see
    :func:`object_point_cloud`'s identical reasoning for its own mask resize.
    """
    mask = np.asarray(mask, dtype=bool)
    if mask.shape == (height, width):
        return mask
    resized = cv2.resize(
        mask.astype(np.uint8), (width, height), interpolation=cv2.INTER_NEAREST
    )
    return resized.astype(bool)


def _mask_characteristic_size_m(mask: np.ndarray, camera: Camera, median_depth: float) -> float:
    """The mask's tight pixel bbox, converted to a single metric length at ``median_depth``.

    Pinhole similar-triangles: a pixel span ``p`` at focal length ``f`` and
    depth ``d`` subtends ``p / f * d`` metres along that axis, exactly (no
    dependence on the noisy per-pixel depth values that corrupt the point
    cloud's own spatial extent -- only the mask's pixel footprint and one
    robust scalar depth are used). The two in-plane axes are then averaged
    into a single characteristic length: which mesh axis is "width" vs
    "height" is not something this function (or its caller, before orientation
    is resolved) can know, so a single orientation-agnostic scalar is what can
    honestly be reported.

    Args:
        mask: ``(H, W)`` boolean object mask, at ``camera``'s resolution.
        camera: Camera whose intrinsics are valid at ``mask``'s resolution.
        median_depth: Robust (median) camera-frame depth of the object, metres.

    Returns:
        The characteristic size, metres.

    Raises:
        ValueError: If ``mask`` has no foreground pixels.
    """
    ys, xs = np.nonzero(mask)
    if ys.size == 0:
        raise ValueError("mask has no foreground pixels")
    bbox_w_px = float(xs.max() - xs.min() + 1)
    bbox_h_px = float(ys.max() - ys.min() + 1)
    width_m = bbox_w_px / camera.K.fx * median_depth
    height_m = bbox_h_px / camera.K.fy * median_depth
    return (width_m + height_m) / 2.0


def _mesh_characteristic_size(mesh_extent: np.ndarray) -> float:
    """A mesh's own canonical size, comparable to :func:`_mask_characteristic_size_m`.

    Averages the two *largest* of the mesh's 3 per-axis extents and ignores
    the smallest -- that smallest one is typically a thickness axis that a
    single 2D mask silhouette never observes anyway, so including it would
    bias the normaliser for exactly the objects (flat-ish ones) where the
    mismatch matters most.
    """
    largest_two = np.sort(np.asarray(mesh_extent, dtype=np.float64))[-2:]
    return float(np.mean(largest_two))


def scale_from_mask(
    mask: np.ndarray, camera: Camera, median_depth: float, mesh: ObjectMesh
) -> float:
    """Recover a mesh's world scale from the mask's pixel footprint, not the cloud's spread.

    Why this exists (measured, not hypothesised -- see the module docstring):
    for a 2x4 LEGO brick (true size ~0.032 x 0.016 x 0.019 m) at 0.63 m, the
    depth map's own noise (MAD 0.0265 m) is *larger* than the object, so the
    back-projected point cloud's spatial extent (0.215 x 0.185 x 0.121 m, even
    after +/-1cm-around-the-median trimming: 0.052 x 0.091 x 0.056 m) cannot
    recover the true size at any filtering strength -- the noise dominates the
    signal, full stop. The mask, in contrast, was segmented perfectly (1664
    px, a tight 82x81 px bbox): its *angular* extent, combined with the
    (reliable) median depth, does recover the size, because it never touches
    the noisy per-pixel depth values at all. See
    :func:`_mask_characteristic_size_m` for the conversion.

    Args:
        mask: ``(H, W)`` boolean object mask, at ``camera``'s resolution.
        camera: Camera whose intrinsics are valid at ``mask``'s resolution.
        median_depth: Robust (median) camera-frame depth of the object,
            metres -- see :func:`object_point_cloud`'s outlier rejection; the
            *median* depth stays reliable even when the depth *spread* does
            not.
        mesh: Canonical-frame mesh to scale. Its own extent is arbitrary (a
            fitted proxy is already roughly metric, but a generator's raw
            output need not be), so the returned factor is normalised against
            the mesh's own characteristic size, not against 1.

    Returns:
        A uniform scale factor ``s`` such that ``s`` times the mesh's own
        characteristic size matches the mask's angular footprint at
        ``median_depth``.

    Raises:
        ValueError: If ``mask`` has no foreground pixels, its shape disagrees
            with ``camera``'s resolution, or ``mesh`` has near-zero extent.
    """
    mask = np.asarray(mask, dtype=bool)
    if mask.shape != (camera.K.height, camera.K.width):
        raise ValueError(
            f"scale_from_mask: mask shape {mask.shape} != camera resolution "
            f"{(camera.K.height, camera.K.width)}"
        )
    mask_size_m = _mask_characteristic_size_m(mask, camera, median_depth)
    mesh_size = _mesh_characteristic_size(np.asarray(mesh.extent, dtype=np.float64))
    if mesh_size < 1e-9:
        raise ValueError("scale_from_mask: mesh has degenerate (near-zero) extent")
    return mask_size_m / mesh_size


def fit_mesh_to_points(
    mesh: ObjectMesh,
    points_world: np.ndarray,
    allow_scale: bool = True,
    refine_icp: bool = True,
    max_iterations: int = 30,
    tol: float = 1e-9,
    n_surface_samples: int = 3000,
    rng: np.random.Generator | None = None,
    view_direction_world: np.ndarray | None = None,
    mask_full: np.ndarray | None = None,
    camera: Camera | None = None,
    orientation_extent_ratio: float = 3.0,
) -> ObjectAlignment:
    """Fit ``mesh`` (canonical frame) into the world frame against ``points_world``.

    Position, scale, and orientation are three different questions, and this
    function answers each from whichever measurement actually supports it
    rather than deriving all three from one (possibly noise-dominated) fit:

    * **Position** comes from the fit below when orientation is trustworthy
      (ICP/Umeyama's correspondence-based translation corrects for the
      single-view-shell bias described below, which a plain centroid cannot).
      When it is not (see below), no correspondence model is trusted either,
      so position falls back to the plain per-axis *median* of
      ``points_world`` -- robust to the same outliers :func:`object_point_cloud`
      already filters, and still meaningful even when the cloud's *shape*
      is not.
    * **Scale** comes from :func:`scale_from_mask` (``scale_source =
      "mask_angular_extent"``) whenever ``mask_full`` and ``camera`` are both
      given, and from the point cloud's own PCA-extent ratio otherwise
      (``scale_source = "depth_extent"``, the original method). The mask path
      exists because of a measured failure of the depth-extent path: for a 2x4
      LEGO brick (true size ~0.032 x 0.016 x 0.019 m) at 0.63 m, depth-map
      noise alone (MAD 0.0265 m) exceeds the object, so the cloud's spatial
      extent (0.215 x 0.185 x 0.121 m) is wrong by up to 7x and no outlier
      filtering recovers the true size from it -- the noise *is* the signal at
      that scale. The mask was segmented correctly (tight 82x81 px bbox), and
      its angular extent at the median depth is accurate to the centimetre.
      ``depth_extent`` remains the default (``mask_full``/``camera`` omitted)
      because it needs no mask/camera and is fine for objects large enough
      that depth noise is a small fraction of their size -- which is most
      objects other than centimetre-scale ones like this.
    * **Orientation** is the one question the point cloud genuinely has to
      answer, and for a noise-dominated cloud it cannot: initialisation is
      PCA/Umeyama (mesh axes matched to the cloud's), and (when
      ``refine_icp``) refined by multi-start ICP -- see the symmetry note
      below -- but only when the cloud shows real 3D structure. Orientation is
      judged **not** confident (``orientation_confident=False`` on the
      returned :class:`~fpgm.objects.types.ObjectAlignment`, ICP skipped, an
      *identity* rotation returned instead) when either: (a) ``mask_full``/
      ``camera`` are given and the cloud's largest PCA extent exceeds
      ``orientation_extent_ratio`` times the mask-derived size (the measured
      LEGO-brick case: ~7x, well past the default 3x -- a cloud that much
      bigger than the object is depth noise, not shape), or (b) the cloud's
      smallest PCA eigenvalue is negligible relative to its largest (fewer
      than 3 well-defined directions -- e.g. a near-planar shell). A
      wrong-but-confident orientation is exactly the failure this function
      must not produce silently, so this case returns a pose that is honestly
      unrotated rather than one that looks fitted but is not.

    PCA alone cannot resolve axis handedness from second moments -- a box's
    covariance is identical under any of its four proper rotational
    symmetries -- so, in the confident case, all four sign-consistent
    candidate rotations are tried, and (when ``refine_icp``) each is refined
    by ICP independently, keeping whichever converges to the lowest final
    RMSE. A single nearest-neighbour pass *before* refinement was tried first
    and isn't discriminating enough on its own: two 170-degree-apart
    candidates can score within a percent of each other pre-refinement, which
    is enough to pick the wrong one and get ICP stuck in a
    locally-consistent but globally wrong basin. Refining every candidate is
    only ~4x one ICP run and is cheap at these point counts.

    Single-view shell: ``points_world`` (see :func:`object_point_cloud`) only
    covers the camera-facing side of the object, not its whole surface. Naive
    ICP against the *whole closed* mesh surface is not safe in general -- it
    was tried and empirically fails for exactly the shape stage 2 produces
    most often: a box has four proper rotational symmetries, and a partial
    single-face-cluster cloud fits an incorrectly-rotated copy of a *symmetric*
    mesh (near-)equally well, because the mesh itself presents an equivalent
    corner elsewhere on its own surface. Restricting correspondences to
    genuinely camera-facing mesh surface is therefore not optional cleverness
    but the fix: pass ``view_direction_world`` (unit vector, object ->
    camera, e.g. ``camera.cam_to_world(np.zeros(3)) - points_world.mean(0)``,
    normalised) whenever a camera is available -- which it always is in the
    real pipeline, since :func:`object_point_cloud` already required one. Both
    the sign-candidate search and every ICP iteration then discard mesh
    samples whose current-orientation outward normal faces away from the
    camera before finding correspondences. Without it (``None``, the default)
    every mesh sample is eligible, which only recovers the correct pose
    reliably for a mesh with no rotational symmetry (e.g. a convex hull of a
    lopsided object) -- pass it for anything box-shaped.

    Args:
        mesh: Canonical-frame mesh to place, e.g. from :mod:`fpgm.objects.proxy`.
        points_world: ``(N, 3)`` observed object points, world frame.
        allow_scale: Fit a single isotropic scale factor; ``False`` pins scale
            to 1.0 (use when the mesh is already known to be metric;
            ``scale_source`` is then ``"fixed"``).
        refine_icp: Run point-to-nearest-visible-surface-point ICP after the
            PCA/Umeyama initialisation, when orientation is confident.
        max_iterations: ICP iteration cap.
        tol: Stop ICP early once both the rotation-matrix Frobenius delta and
            the translation delta (metres) drop below this in one iteration.
        n_surface_samples: Points sampled uniformly over the mesh surface to
            build the ICP correspondence target.
        rng: Surface-sampling RNG; a fresh ``default_rng()`` is used if omitted.
        view_direction_world: Unit vector, object -> camera, in world frame.
            See above; strongly recommended for any mesh with rotational
            symmetry.
        mask_full: ``(H, W)`` boolean object mask (any resolution; resized
            nearest-neighbour to ``camera``'s if needed). Enables the
            mask-derived scale and the extent-ratio orientation-confidence
            check above; must be given together with ``camera``.
        camera: Camera whose intrinsics/extrinsic apply to ``points_world`` and
            (once resized) ``mask_full``. Must be given together with
            ``mask_full``.
        orientation_extent_ratio: How many times larger than the mask-derived
            size the cloud's largest PCA extent may be before orientation is
            judged unrecoverable. Default 3x is comfortably below the ~7x
            measured on the noise-dominated LEGO-brick case, and comfortably
            above ordinary fit slop for a well-resolved object.

    Returns:
        The fitted :class:`~fpgm.objects.types.ObjectAlignment`.

    Raises:
        AlignmentError: Too few points; the cloud is near-collinear (PCA
            orientation undefined, not even a fallback pose is meaningful); or
            exactly one of ``mask_full``/``camera`` was given.
    """
    points_world = np.asarray(points_world, dtype=np.float64)
    if points_world.ndim != 2 or points_world.shape[1] != 3:
        raise AlignmentError(
            f"fit_mesh_to_points: expected (N, 3) points, got {points_world.shape}"
        )
    if points_world.shape[0] < _MIN_FIT_POINTS:
        raise AlignmentError(
            f"fit_mesh_to_points: need >= {_MIN_FIT_POINTS} points, got "
            f"{points_world.shape[0]}"
        )
    if (mask_full is None) != (camera is None):
        raise AlignmentError(
            "fit_mesh_to_points: mask_full and camera must be given together (or both omitted)"
        )
    if view_direction_world is not None:
        view_direction_world = np.asarray(view_direction_world, dtype=np.float64)
        view_direction_world = view_direction_world / np.linalg.norm(view_direction_world)

    position_median = np.median(points_world, axis=0)

    cloud_centroid, cloud_axes, cloud_eigvals = _pca(points_world)
    if cloud_eigvals[1] <= _COLLINEAR_EIGVAL_RATIO * cloud_eigvals[0]:
        raise AlignmentError(
            "fit_mesh_to_points: point cloud is near-collinear (PCA eigenvalues "
            f"{cloud_eigvals.tolist()}); orientation is undefined"
        )
    cloud_local = points_world @ cloud_axes
    cloud_extent = np.maximum(cloud_local.max(axis=0) - cloud_local.min(axis=0), 1e-9)

    mesh_vertices = np.asarray(mesh.vertices, dtype=np.float64)
    mesh_centroid, mesh_axes, _ = _pca(mesh_vertices)
    mesh_local = mesh_vertices @ mesh_axes
    mesh_extent = np.maximum(mesh_local.max(axis=0) - mesh_local.min(axis=0), 1e-9)

    # --- scale: mask angular extent when available, depth extent otherwise ---
    mask_size_m: float | None = None
    if mask_full is not None and camera is not None:
        mask_at_res = _resize_mask_nearest(mask_full, camera.K.height, camera.K.width)
        median_depth = float(np.median(camera.world_to_cam(points_world)[..., 2]))
        mask_size_m = _mask_characteristic_size_m(mask_at_res, camera, median_depth)

    if not allow_scale:
        scale = 1.0
        scale_source = "fixed"
    elif mask_size_m is not None:
        mesh_size = _mesh_characteristic_size(mesh_extent)
        if mesh_size < 1e-9:
            raise AlignmentError("fit_mesh_to_points: mesh has degenerate (near-zero) extent")
        scale = mask_size_m / mesh_size
        scale_source = "mask_angular_extent"
    else:
        scale = float(np.mean(cloud_extent / mesh_extent))
        scale_source = "depth_extent"

    # --- orientation confidence: is there any structure worth fitting? ---
    notes: list[str] = []
    orientation_confident = True
    if cloud_eigvals[2] <= _LOW_STRUCTURE_EIGVAL_RATIO * cloud_eigvals[0]:
        orientation_confident = False
        notes.append(
            f"low structure: smallest/largest PCA eigenvalue ratio "
            f"{cloud_eigvals[2] / cloud_eigvals[0]:.2e} <= {_LOW_STRUCTURE_EIGVAL_RATIO:.0e}"
        )
    if mask_size_m is not None:
        cloud_extent_max = float(np.max(cloud_extent))
        if cloud_extent_max > orientation_extent_ratio * mask_size_m:
            orientation_confident = False
            notes.append(
                f"cloud extent {cloud_extent_max:.4f}m exceeds "
                f"{orientation_extent_ratio:g}x the mask-derived size {mask_size_m:.4f}m "
                "(depth noise likely dominates the cloud's shape)"
            )

    if not orientation_confident:
        # No correspondence model is trusted here (that is exactly what
        # "not confident" means), so this is a placement, not a fit: identity
        # rotation, translation from the robust median position. rmse/
        # inlier_fraction are still reported as plain diagnostics (how far the
        # placed-but-unrotated mesh sits from the cloud), not as evidence the
        # orientation is right.
        rotation = np.eye(3, dtype=np.float64)
        translation = position_median - scale * mesh_centroid
        # Raw (uncentered) vertices, matching how `transform` is actually
        # applied everywhere else (e.g. save_debug): world = scale*v + t.
        transformed = scale * mesh_vertices + translation
        dists, _ = cKDTree(transformed).query(points_world)
        rmse = float(np.sqrt(np.mean(dists**2)))
        inlier_fraction = _inlier_fraction(dists, cloud_extent)
        notes.append(f"orientation fallback: identity rotation, scale_source={scale_source}")
    else:
        tri = trimesh.Trimesh(vertices=mesh_vertices, faces=np.asarray(mesh.faces), process=False)
        rng = rng if rng is not None else np.random.default_rng()
        samples_canonical, face_idx = trimesh.sample.sample_surface(
            tri, n_surface_samples, seed=rng
        )
        samples_canonical = np.asarray(samples_canonical, dtype=np.float64)
        sample_normals_canonical = np.asarray(tri.face_normals[face_idx], dtype=np.float64)

        # Multi-start: refine *every* one of the 4 sign candidates and keep the
        # lowest final RMSE, rather than picking one candidate up front from a
        # single nearest-neighbour pass and only refining that. A one-shot
        # residual before refinement was tried and is not discriminating enough
        # -- two candidates 170 degrees apart can score within a percent of
        # each other pre-ICP, and picking the marginally-better-looking wrong
        # one before refinement gets ICP stuck in its (locally consistent,
        # globally wrong) basin. Running refinement per candidate is only ~4x
        # the cost of one ICP run, cheap at this point count.
        best: tuple[np.ndarray, float, np.ndarray, float, float] | None = None
        for signs in _SIGN_CANDIDATES:
            r0 = cloud_axes @ np.diag(signs) @ mesh_axes.T
            t0 = cloud_centroid - scale * (r0 @ mesh_centroid)
            if refine_icp:
                candidate = _icp_refine(
                    samples_canonical,
                    sample_normals_canonical,
                    points_world,
                    r0,
                    scale,
                    t0,
                    allow_scale=allow_scale,
                    max_iterations=max_iterations,
                    tol=tol,
                    view_direction_world=view_direction_world,
                )
            else:
                subset = _visible_subset(
                    samples_canonical, sample_normals_canonical, r0, view_direction_world
                )
                transformed = (scale * (r0 @ (subset - mesh_centroid).T)).T + t0
                dists, _ = cKDTree(transformed).query(points_world)
                rmse = float(np.sqrt(np.mean(dists**2)))
                candidate = (r0, scale, t0, rmse, _inlier_fraction(dists, cloud_extent))
            if best is None or candidate[3] < best[3]:
                best = candidate
        assert best is not None  # _SIGN_CANDIDATES is non-empty
        rotation, scale, translation, rmse, inlier_fraction = best
        notes.append(
            f"orientation fitted via {'ICP-refined' if refine_icp else 'PCA-only'} "
            f"multi-start, scale_source={scale_source}"
        )

    transform = np.eye(4, dtype=np.float64)
    transform[:3, :3] = scale * rotation
    transform[:3, 3] = translation

    return ObjectAlignment(
        transform=transform,
        scale=scale,
        points_world=points_world,
        rmse_m=rmse,
        inlier_fraction=inlier_fraction,
        orientation_confident=orientation_confident,
        scale_source=scale_source,
        notes="; ".join(notes),
    )


def _pca(points: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Centroid, right-handed PCA axes (columns, descending variance), eigenvalues."""
    centroid = points.mean(axis=0)
    centered = points - centroid
    cov = (centered.T @ centered) / points.shape[0]
    eigvals, eigvecs = np.linalg.eigh(cov)  # ascending
    order = np.argsort(eigvals)[::-1]
    eigvals, axes = eigvals[order], eigvecs[:, order]
    if np.linalg.det(axes) < 0:
        axes[:, -1] *= -1.0
    return centroid, axes, eigvals


def _umeyama(
    src: np.ndarray, dst: np.ndarray, with_scale: bool
) -> tuple[np.ndarray, float, np.ndarray]:
    """Closed-form similarity transform minimising ``sum ||s*R@src_i + t - dst_i||^2``.

    Umeyama (1991). ``src``/``dst`` must be corresponding ``(N, 3)`` point sets.

    Returns:
        ``(R, s, t)``: a ``(3, 3)`` proper rotation, a scalar scale, and a
        ``(3,)`` translation.
    """
    n = src.shape[0]
    mu_src, mu_dst = src.mean(axis=0), dst.mean(axis=0)
    src_c, dst_c = src - mu_src, dst - mu_dst

    cov = (dst_c.T @ src_c) / n
    u, s_vals, vt = np.linalg.svd(cov)
    d = np.ones(3)
    if np.linalg.det(u) * np.linalg.det(vt) < 0:
        d[-1] = -1.0
    r = u @ np.diag(d) @ vt

    if with_scale:
        var_src = float((src_c**2).sum(axis=1).mean())
        scale = float(np.sum(s_vals * d) / var_src) if var_src > 1e-15 else 1.0
    else:
        scale = 1.0

    t = mu_dst - scale * (r @ mu_src)
    return r, scale, t


#: Below this many visible samples, visibility filtering is treated as
#: degenerate (e.g. a pathological view direction) and the full sample set is
#: used instead, rather than fitting against a near-empty target.
_MIN_VISIBLE_SAMPLES = 20


def _visible_subset(
    samples_canonical: np.ndarray,
    sample_normals_canonical: np.ndarray,
    rotation: np.ndarray,
    view_direction_world: np.ndarray | None,
) -> np.ndarray:
    """Mesh surface samples whose current-orientation normal faces the camera.

    A pure rotation (no translation/scale) carries a normal correctly, so only
    ``rotation`` is applied. Returns all of ``samples_canonical`` unchanged if
    ``view_direction_world`` is ``None`` or too few samples would survive.
    """
    if view_direction_world is None:
        return samples_canonical
    normals_world = (rotation @ sample_normals_canonical.T).T
    visible = normals_world @ view_direction_world > 0.0
    if visible.sum() < _MIN_VISIBLE_SAMPLES:
        return samples_canonical
    return samples_canonical[visible]


def _icp_refine(
    samples_canonical: np.ndarray,
    sample_normals_canonical: np.ndarray,
    points_world: np.ndarray,
    rotation: np.ndarray,
    scale: float,
    translation: np.ndarray,
    allow_scale: bool,
    max_iterations: int,
    tol: float,
    view_direction_world: np.ndarray | None,
) -> tuple[np.ndarray, float, np.ndarray, float, float]:
    """Point-to-nearest-visible-surface-point ICP.

    Correspondences are found by bringing the (world-frame) cloud into the
    mesh's static canonical frame each iteration and querying a KD-tree over
    the currently camera-facing samples (see :func:`_visible_subset`), rather
    than rebuilding a KD-tree on the *transformed* mesh every iteration -- the
    canonical-frame sample positions never move, only which of them are
    eligible does.
    """
    cloud_extent = points_world.max(axis=0) - points_world.min(axis=0)

    for _ in range(max_iterations):
        subset = _visible_subset(
            samples_canonical, sample_normals_canonical, rotation, view_direction_world
        )
        local_query = ((points_world - translation) @ rotation) / scale
        _, idx = cKDTree(subset).query(local_query)
        correspondences = subset[idx]

        new_rotation, new_scale, new_translation = _umeyama(
            correspondences, points_world, with_scale=allow_scale
        )
        if not allow_scale:
            new_scale = scale

        rot_delta = float(np.linalg.norm(new_rotation - rotation))
        trans_delta = float(np.linalg.norm(new_translation - translation))
        rotation, scale, translation = new_rotation, new_scale, new_translation
        if rot_delta < tol and trans_delta < tol:
            break

    final_subset = _visible_subset(
        samples_canonical, sample_normals_canonical, rotation, view_direction_world
    )
    transformed = (scale * (rotation @ final_subset.T)).T + translation
    dists, _ = cKDTree(transformed).query(points_world)
    rmse = float(np.sqrt(np.mean(dists**2)))
    inlier_fraction = _inlier_fraction(dists, cloud_extent)
    return rotation, scale, translation, rmse, inlier_fraction


def _inlier_fraction(dists: np.ndarray, cloud_extent: np.ndarray, frac: float = 0.1) -> float:
    """Fraction of correspondences within ``frac`` of the cloud's largest extent.

    Scale-relative rather than an absolute distance cutoff, so the same
    default works for a 3 cm brick and a 30 cm box alike.
    """
    threshold = max(frac * float(np.max(cloud_extent)), 1e-6)
    return float(np.mean(dists <= threshold))


def save_debug(
    alignment: ObjectAlignment,
    camera: Camera,
    frame_bgr: np.ndarray,
    out_dir: Path,
    mesh: ObjectMesh | None = None,
) -> StageArtifacts:
    """Reproject the fitted mesh silhouette and observed cloud onto ``frame_bgr``.

    Args:
        alignment: Stage 3 output to visualise.
        camera: Camera the alignment's world points/mesh should be seen through.
        frame_bgr: The frame to draw over (BGR, uint8).
        out_dir: Directory to write into (created if missing).
        mesh: The mesh that was fitted, in its own canonical frame, if
            available -- :class:`~fpgm.objects.types.ObjectAlignment` itself
            only carries the fitted transform, not the mesh. When given, its
            vertices are transformed and projected to draw a silhouette; when
            omitted, the observed point cloud's own convex hull stands in,
            which is still an informative "does the fit look plausible" check.

    Returns:
        A :class:`~fpgm.objects.types.StageArtifacts` pointing at the overlay
        PNG and a stats JSON.
    """
    out_dir = ensure_dir(Path(out_dir))
    artifacts = StageArtifacts(stage="align", directory=out_dir)

    overlay = frame_bgr.copy()
    uv_points, depth_points = camera.project(alignment.points_world)
    for u, v in np.round(uv_points[depth_points > 0]).astype(int):
        cv2.circle(overlay, (int(u), int(v)), 2, (0, 215, 255), -1, lineType=cv2.LINE_AA)

    if mesh is not None:
        verts_world = (alignment.transform[:3, :3] @ mesh.vertices.T).T + alignment.transform[:3, 3]
    else:
        verts_world = alignment.points_world
    uv_mesh, depth_mesh = camera.project(verts_world)
    visible = uv_mesh[depth_mesh > 0]
    if visible.shape[0] >= 3:
        hull = cv2.convexHull(np.round(visible).astype(np.int32))
        cv2.polylines(overlay, [hull], True, color_for(1), 2, lineType=cv2.LINE_AA)

    overlay_path = out_dir / f"align_overlay_{alignment.frame_idx:06d}.png"
    cv2.imwrite(str(overlay_path), overlay)
    artifacts.add("overlay", overlay_path)

    extent = alignment.points_world.max(axis=0) - alignment.points_world.min(axis=0)
    stats = {
        "frame_idx": alignment.frame_idx,
        "n_points": int(alignment.points_world.shape[0]),
        "extent_m": extent.tolist(),
        "scale": alignment.scale,
        "rmse_m": alignment.rmse_m,
        "inlier_fraction": alignment.inlier_fraction,
        "position_m": alignment.position.tolist(),
    }
    stats_path = out_dir / f"align_stats_{alignment.frame_idx:06d}.json"
    stats_path.write_text(json.dumps(stats, indent=2))
    artifacts.add("stats", stats_path)
    artifacts.stats = stats
    return artifacts