src/fpgm/objects/proxy.py

"""Stage 2 stand-in: proxy meshes fitted to the observed point cloud.

The real generator (TRELLIS.2, driven off :mod:`fpgm.objects.crop`'s output) is
not always available -- this module lets stages 3-5 run end-to-end without it,
by fitting a closed proxy mesh directly to the stage-3 point cloud. Both proxies
are returned in their own *canonical* frame (centred near the origin, axes =
the point cloud's PCA axes) with the fitted world placement stashed in
``ObjectMesh.metadata["world_transform"]`` -- this matches the convention
:mod:`fpgm.objects.align` expects a mesh to already be in, so a proxy and a
future real generator output are interchangeable inputs to that stage.

Both builders take an optional ``point_colors`` (the same per-point colours
:func:`fpgm.objects.align.object_point_cloud` returns) and, when given, set
``ObjectMesh.vertex_colors`` from them -- each mesh vertex takes its *nearest*
observed point's colour (see :func:`_nearest_point_colors`). Without this a
proxy mesh has no colour at all and renders as a flat grey blob regardless of
the object's real appearance.
"""

from __future__ import annotations

import numpy as np
import trimesh
from scipy.spatial import ConvexHull, QhullError, cKDTree

from fpgm.objects.types import MeshSource, ObjectMesh
from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

#: PCA needs at least this many points to define three axes non-degenerately.
_MIN_POINTS = 4
#: Guard against a zero-thickness box dimension (e.g. an almost-planar cloud).
_MIN_EXTENT_M = 1e-6


def _nearest_point_colors(
    verts_world: np.ndarray, points_world: np.ndarray, point_colors: np.ndarray
) -> np.ndarray:
    """Per-vertex colour: the *nearest observed point's* colour, in world frame.

    Picked over a single flat (e.g. median) colour for the whole mesh because
    a real object is rarely one colour -- a nearest-point lookup lets
    different faces of the proxy pick up whatever was actually observed near
    them (e.g. a brick's printed top vs its plain sides), while still never
    inventing a colour that wasn't observed.

    Args:
        verts_world: ``(V, 3)`` mesh vertices, world frame.
        points_world: ``(N, 3)`` observed object points, world frame.
        point_colors: ``(N, 3)`` uint8 colours, one per ``points_world`` row.

    Returns:
        ``(V, 3)`` uint8 colours, one per vertex.
    """
    tree = cKDTree(points_world)
    _, idx = tree.query(verts_world)
    return np.asarray(point_colors, dtype=np.uint8)[idx]


def _pca_frame(points_world: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Centroid, right-handed PCA axes (columns, descending variance), and points in that frame.

    Args:
        points_world: ``(N, 3)`` observed points.

    Returns:
        ``(centroid, axes, local)``: ``centroid`` is ``(3,)``, ``axes`` is
        ``(3, 3)`` with columns the principal directions (largest variance
        first) forming a proper rotation, and ``local`` is
        ``(points_world - centroid) @ axes``.

    Raises:
        ValueError: If there are fewer than :data:`_MIN_POINTS` points.
    """
    points_world = np.asarray(points_world, dtype=np.float64)
    if points_world.ndim != 2 or points_world.shape[1] != 3:
        raise ValueError(f"expected (N, 3) points, got shape {points_world.shape}")
    if points_world.shape[0] < _MIN_POINTS:
        raise ValueError(
            f"need at least {_MIN_POINTS} points to fit a proxy mesh, got "
            f"{points_world.shape[0]}"
        )

    centroid = points_world.mean(axis=0)
    centered = points_world - centroid
    cov = (centered.T @ centered) / points_world.shape[0]
    eigvals, eigvecs = np.linalg.eigh(cov)  # ascending
    order = np.argsort(eigvals)[::-1]
    axes = eigvecs[:, order]
    if np.linalg.det(axes) < 0:
        axes[:, -1] *= -1.0  # flip the least-informative axis to keep a proper rotation

    local = centered @ axes
    return centroid, axes, local


def oriented_box_mesh(
    points_world: np.ndarray, point_colors: np.ndarray | None = None
) -> ObjectMesh:
    """PCA-oriented bounding box of the observed points.

    The box is built axis-aligned in its own canonical frame (extents = the
    points' range along each PCA axis), so it is the tightest box that
    contains every observed point under *some* rotation -- not the (looser)
    world-axis-aligned bounding box.

    Args:
        points_world: ``(N, 3)`` observed object points, world frame.
        point_colors: ``(N, 3)`` uint8 colours, one per ``points_world`` row,
            e.g. from :func:`fpgm.objects.align.object_point_cloud`. When
            given, each box vertex gets the colour of its nearest observed
            point (see :func:`_nearest_point_colors`); when omitted (the
            default), ``vertex_colors`` is ``None`` and the mesh is
            colourless, as before.

    Returns:
        A closed, watertight :class:`~fpgm.objects.types.ObjectMesh` with
        ``source=MeshSource.PROXY_BOX`` and
        ``metadata["world_transform"]`` (a ``(4, 4)`` matrix, no scale --
        the box is already metric) mapping its canonical vertices to world
        coordinates.

    Raises:
        ValueError: If there are fewer than :data:`_MIN_POINTS` points.
    """
    centroid, axes, local = _pca_frame(points_world)
    lo, hi = local.min(axis=0), local.max(axis=0)
    extents = np.maximum(hi - lo, _MIN_EXTENT_M)
    local_center = (hi + lo) / 2.0

    box = trimesh.creation.box(extents=extents)
    box_vertices = np.asarray(box.vertices, dtype=np.float64)

    world_transform = np.eye(4, dtype=np.float64)
    world_transform[:3, :3] = axes
    # The box's own vertices are centred on the *bbox* centre, not the point
    # centroid the local frame is centred on (a skewed cloud can have those
    # differ), so the extra local_center offset is needed here.
    world_transform[:3, 3] = centroid + axes @ local_center

    vertex_colors = None
    if point_colors is not None:
        verts_world = (world_transform[:3, :3] @ box_vertices.T).T + world_transform[:3, 3]
        vertex_colors = _nearest_point_colors(verts_world, points_world, point_colors)

    return ObjectMesh(
        vertices=box_vertices,
        faces=np.asarray(box.faces, dtype=np.int64),
        source=MeshSource.PROXY_BOX,
        vertex_colors=vertex_colors,
        metadata={
            "world_transform": world_transform,
            "half_extents": (extents / 2.0).tolist(),
            "n_points": int(local.shape[0]),
        },
    )


def convex_hull_mesh(
    points_world: np.ndarray, point_colors: np.ndarray | None = None
) -> ObjectMesh:
    """Convex hull of the observed points, closer to the true silhouette than a box.

    Args:
        points_world: ``(N, 3)`` observed object points, world frame.
        point_colors: ``(N, 3)`` uint8 colours, one per ``points_world`` row,
            e.g. from :func:`fpgm.objects.align.object_point_cloud`. When
            given, each hull vertex gets the colour of its nearest observed
            point (see :func:`_nearest_point_colors`); when omitted (the
            default), ``vertex_colors`` is ``None`` and the mesh is
            colourless, as before.

    Returns:
        A closed, watertight :class:`~fpgm.objects.types.ObjectMesh` with
        ``source=MeshSource.PROXY_HULL`` and ``metadata["world_transform"]``
        mapping its canonical vertices to world coordinates.

    Raises:
        ValueError: If there are fewer than :data:`_MIN_POINTS` points, the
            points are degenerate (coplanar/collinear -- no 3D hull exists),
            or the resulting hull is not watertight.
    """
    centroid, axes, local = _pca_frame(points_world)
    try:
        hull = ConvexHull(local)
    except QhullError as exc:
        raise ValueError(
            "convex_hull_mesh: points are degenerate (coplanar/collinear); "
            "cannot compute a 3D convex hull"
        ) from exc

    # ConvexHull.simplices indexes into the full input array; most of those
    # points are interior and never referenced, so compact down to only the
    # vertices actually on the hull before building the mesh.
    used = np.unique(hull.simplices)
    remap = np.full(local.shape[0], -1, dtype=np.int64)
    remap[used] = np.arange(used.shape[0])
    vertices = local[used]
    faces = remap[hull.simplices]

    tri = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
    tri.fix_normals()  # qhull's facet winding isn't guaranteed outward-consistent
    if not tri.is_watertight:
        raise ValueError("convex_hull_mesh: resulting hull mesh is not watertight")

    # `local` is already centred at `centroid` (see _pca_frame), so unlike the
    # box no extra offset is needed: canonical -> world is just axes + centroid.
    world_transform = np.eye(4, dtype=np.float64)
    world_transform[:3, :3] = axes
    world_transform[:3, 3] = centroid

    hull_vertices = np.asarray(tri.vertices, dtype=np.float64)
    vertex_colors = None
    if point_colors is not None:
        verts_world = (world_transform[:3, :3] @ hull_vertices.T).T + world_transform[:3, 3]
        vertex_colors = _nearest_point_colors(verts_world, points_world, point_colors)

    return ObjectMesh(
        vertices=hull_vertices,
        faces=np.asarray(tri.faces, dtype=np.int64),
        source=MeshSource.PROXY_HULL,
        vertex_colors=vertex_colors,
        metadata={"world_transform": world_transform, "n_points": int(local.shape[0])},
    )