tests/test_object_align.py

"""Synthetic, data-free tests for fpgm.objects.align -- no GPU, no data files."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import pytest
import trimesh

from fpgm.geometry.camera import Camera
from fpgm.geometry.transforms import rotvec_to_matrix
from fpgm.objects.align import (
    AlignmentError,
    _pca,
    fit_mesh_to_points,
    object_point_cloud,
    scale_from_mask,
)
from fpgm.objects.proxy import convex_hull_mesh, oriented_box_mesh
from fpgm.objects.types import MeshSource, ObjectMesh
from fpgm.types import CameraIntrinsics


@dataclass
class _DepthClip:
    """Minimal stand-in for align.DepthRgbSource."""

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


def _make_camera(width: int, height: int, fx: float = 300.0, fy: float = 300.0) -> Camera:
    """A camera at the world origin looking straight down +Z (identity extrinsic)."""
    intrinsics = CameraIntrinsics(
        fx=fx, fy=fy, cx=width / 2.0, cy=height / 2.0, width=width, height=height
    )
    return Camera(intrinsics, np.eye(4))


def _rescale_to_mad(raw: np.ndarray, target_mad: float) -> np.ndarray:
    """Rescale zero-median ``raw`` so its median absolute deviation is exactly ``target_mad``.

    Lets a test specify a *measured* MAD directly (e.g. the 0.0265 m figure
    from the demo LEGO brick) without depending on a particular sample count
    for a random draw to converge to it.
    """
    mad = float(np.median(np.abs(raw - np.median(raw))))
    assert mad > 1e-12, "test setup: raw noise must have nonzero spread"
    return raw * (target_mad / mad)


def _front_face_depth_image(
    camera: Camera, depth_shape: tuple[int, int], box_center: np.ndarray, half_extents: np.ndarray
) -> tuple[np.ndarray, tuple[float, float, float, float]]:
    """Depth image of a box's near face only -- what a single depth camera actually sees.

    Returns the ``(H, W)`` uint16-mm depth array and the ground-truth near-face
    footprint ``(x0, y0, x1, y1)`` in world metres, for assertions.
    """
    h, w = depth_shape
    near_z = box_center[2] - half_extents[2]
    x0, x1 = box_center[0] - half_extents[0], box_center[0] + half_extents[0]
    y0, y1 = box_center[1] - half_extents[1], box_center[1] + half_extents[1]

    corners_world = np.array([[x0, y0, near_z], [x1, y1, near_z]])
    uv, _ = camera.project(corners_world)  # fx, fy > 0 and no rotation -> monotonic in x, y
    (u0, v0), (u1, v1) = uv
    px0, px1 = int(round(min(u0, u1))), int(round(max(u0, u1)))
    py0, py1 = int(round(min(v0, v1))), int(round(max(v0, v1)))

    depth_mm = np.zeros((h, w), dtype=np.uint16)
    depth_mm[py0:py1, px0:px1] = int(round(near_z * 1000.0))
    return depth_mm, (x0, y0, x1, y1)


class TestObjectPointCloud:
    def test_recovers_centroid_and_extent_of_known_box(self) -> None:
        depth_shape = (90, 160)  # (H, W), annotation resolution
        camera = _make_camera(width=depth_shape[1], height=depth_shape[0])
        box_center = np.array([0.0, 0.0, 0.6])
        half_extents = np.array([0.016, 0.008, 0.0095])  # ~ LEGO-brick scale

        depth_mm, (x0, y0, x1, y1) = _front_face_depth_image(
            camera, depth_shape, box_center, half_extents
        )

        # mask_full is deliberately at 4x the depth resolution -- mirrors the
        # real 1280x720 video vs 320x180 annotation-depth mismatch -- so this
        # also exercises the nearest-neighbour downscale path.
        full_h, full_w = depth_shape[0] * 4, depth_shape[1] * 4
        mask_full = np.zeros((full_h, full_w), dtype=bool)
        mask_full_small = depth_mm > 0
        ys, xs = np.nonzero(mask_full_small)
        mask_full[ys.min() * 4 : (ys.max() + 1) * 4, xs.min() * 4 : (xs.max() + 1) * 4] = True

        clip = _DepthClip(initial_depth=depth_mm)
        points_world, colors = object_point_cloud(
            clip, camera, mask_full, frame_shape=(full_h, full_w)
        )

        assert points_world.shape[0] > 0
        assert colors.shape == (points_world.shape[0], 3)

        # A single depth image only sees the near face: Z should collapse to
        # (approximately) the near-face depth, while X/Y span the box's footprint.
        near_z = box_center[2] - half_extents[2]
        # depth is quantised to whole millimetres on write, so sub-mm error is expected.
        assert points_world[:, 2] == pytest.approx(near_z, abs=1e-3)

        extent = points_world.max(axis=0) - points_world.min(axis=0)
        true_extent_xy = np.array([x1 - x0, y1 - y0])
        # Pixel-quantisation tolerance: one annotation pixel at fx=300, z=0.6m.
        px_tol = box_center[2] / 300.0
        assert extent[:2] == pytest.approx(true_extent_xy, abs=3 * px_tol)
        assert extent[2] < 1e-6

        centroid = points_world.mean(axis=0)
        true_center_xy = np.array([(x0 + x1) / 2.0, (y0 + y1) / 2.0])
        assert centroid[:2] == pytest.approx(true_center_xy, abs=3 * px_tol)

    def test_outlier_rejection_removes_leaky_background_pixels(self) -> None:
        """Regression test for the 0.39 m bounding box bug from a leaky mask."""
        depth_shape = (90, 160)
        camera = _make_camera(width=depth_shape[1], height=depth_shape[0])
        box_center = np.array([0.0, 0.0, 0.6])
        half_extents = np.array([0.016, 0.008, 0.0095])

        depth_mm, (x0, y0, x1, y1) = _front_face_depth_image(
            camera, depth_shape, box_center, half_extents
        )
        object_mask = depth_mm > 0
        ys, xs = np.nonzero(object_mask)

        # Leak a handful of far-background pixels into the mask, at a depth a
        # metre+ beyond the object -- exactly the scenario that produced the
        # 0.39 m box for a 0.03 m brick.
        leak_depth_mm = int(1.8 * 1000.0)
        leak_ys = np.clip(
            [ys.min() - 3, ys.min() - 3, ys.max() + 3, ys.max() + 3, ys.min() - 3],
            0,
            depth_shape[0] - 1,
        )
        leak_xs = np.clip(
            [xs.min() - 3, xs.max() + 3, xs.min() - 3, xs.max() + 3, xs.min() - 5],
            0,
            depth_shape[1] - 1,
        )
        depth_mm[leak_ys, leak_xs] = leak_depth_mm

        mask_full = object_mask.copy()
        mask_full[leak_ys, leak_xs] = True
        n_leak = int((mask_full & ~object_mask).sum())
        assert n_leak >= 3, "test setup: leaked pixels must land outside the object mask"

        clip = _DepthClip(initial_depth=depth_mm)
        points_world, _ = object_point_cloud(clip, camera, mask_full, frame_shape=depth_shape)

        # None of the leaked far-background depth should survive rejection.
        assert points_world[:, 2].max() < 1.0
        extent = points_world.max(axis=0) - points_world.min(axis=0)
        assert np.all(extent < 0.05)  # true object is a few cm; 0.39 m would fail this


def _box_mesh(extents: tuple[float, float, float]) -> ObjectMesh:
    box = trimesh.creation.box(extents=extents)
    return ObjectMesh(
        vertices=np.asarray(box.vertices, dtype=np.float64),
        faces=np.asarray(box.faces, dtype=np.int64),
        source=MeshSource.PROXY_BOX,
    )


def _lopsided_hull_mesh(rng: np.random.Generator) -> ObjectMesh:
    """A convex hull with no rotational symmetry -- see the note on box symmetry below.

    Built by jittering a box's 8 corners and re-hulling, so it stays roughly
    brick-sized/-shaped (what :mod:`fpgm.objects.proxy` actually produces from
    noisy real points) while breaking the *exact* symmetry a perfect box has.
    """
    box = trimesh.creation.box(extents=(0.032, 0.019, 0.024))
    verts = np.asarray(box.vertices) + rng.uniform(-0.004, 0.004, size=box.vertices.shape)
    hull = trimesh.Trimesh(vertices=verts, process=True).convex_hull
    return ObjectMesh(
        vertices=np.asarray(hull.vertices, dtype=np.float64),
        faces=np.asarray(hull.faces, dtype=np.int64),
        source=MeshSource.PROXY_HULL,
    )


def _sample_visible_faces(
    mesh: ObjectMesh, view_dir: np.ndarray, n: int, rng: np.random.Generator, thresh: float = 0.15
) -> np.ndarray:
    """Sample points only from faces facing ``view_dir`` -- a single-view shell."""
    tri = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, process=False)
    front = tri.face_normals @ view_dir > thresh
    sub = tri.submesh([front], append=True)
    pts, _ = trimesh.sample.sample_surface(sub, n, seed=rng)
    return np.asarray(pts, dtype=np.float64)


def _rotation_angle_error_deg(true_rotation: np.ndarray, recovered_rotation: np.ndarray) -> float:
    r_err = true_rotation.T @ recovered_rotation
    return float(np.degrees(np.arccos(np.clip((np.trace(r_err) - 1.0) / 2.0, -1.0, 1.0))))


class TestFitMeshToPoints:
    """
    A note on why these tests use a *lopsided* hull rather than a plain box:

    a plain rectangular box is exactly, provably invariant under 180-degree
    rotation about any of its 3 principal axes (each such rotation maps the
    box's occupied region onto itself exactly, as a set, regardless of how
    distinct its side lengths are). That is a fact about the box's geometry,
    not a property of the fitting algorithm: matching a single-view partial
    shell against a *plain* box is genuinely 4-fold ambiguous, and no
    registration method can recover "the" orientation from geometry alone,
    because the 4 candidates predict bit-for-bit identical observations. This
    was verified empirically while building this test -- a plain box's 4
    sign-candidates converge to indistinguishable (to 1e-6) final RMSE no
    matter how oblique the view or how distinct the box's extents are.

    A real fitted proxy is never perfectly symmetric like this (it is fit to
    noisy real points), so a mildly irregular hull -- still fit with the same
    machinery, same 4-candidate multi-start ICP -- is the fair and realistic
    stand-in, and does have one genuinely-best answer to recover.
    """

    def test_recovers_known_scale_rotation_translation(self) -> None:
        rng = np.random.default_rng(0)
        mesh = _lopsided_hull_mesh(rng)

        true_scale = 1.0
        true_rotation = rotvec_to_matrix(rng.uniform(-0.8, 0.8, size=3))
        true_translation = np.array([0.1, -0.05, 0.6])

        view_dir = np.array([1.0, 1.0, 1.0]) / np.sqrt(3.0)
        local_points = _sample_visible_faces(mesh, view_dir, 700, rng)
        points_world = (true_scale * (true_rotation @ local_points.T)).T + true_translation

        # In the real pipeline this comes from the camera that produced the
        # cloud, e.g. `camera.cam_to_world(np.zeros(3)) - points_world.mean(0)`,
        # normalised. Here it's the same direction (rotated into world frame)
        # used above to pick which faces were "visible" when sampling.
        view_direction_world = true_rotation @ view_dir

        alignment = fit_mesh_to_points(
            mesh,
            points_world,
            allow_scale=True,
            refine_icp=True,
            rng=np.random.default_rng(100),
            view_direction_world=view_direction_world,
        )

        assert alignment.scale == pytest.approx(true_scale, rel=0.05)
        recovered_rotation = alignment.transform[:3, :3] / alignment.scale
        assert _rotation_angle_error_deg(true_rotation, recovered_rotation) < 5.0
        assert alignment.position == pytest.approx(true_translation, abs=0.005)
        assert alignment.rmse_m < 0.003
        assert alignment.inlier_fraction > 0.8
        # A clean, well-resolved synthetic cloud has real 3D structure -- ICP
        # orientation should be trusted (and no mask was given, so the scale
        # fell back to the original depth-extent estimate).
        assert alignment.orientation_confident is True
        assert alignment.scale_source == "depth_extent"

    def test_recovers_correct_rotation_despite_flipped_pca_sign_init(self) -> None:
        """Regression test: a naive single-sign PCA init lands 180-degrees off here.

        Pins the multi-start-ICP fix directly: this specific (seeded) cloud is
        constructed so that pairing the cloud's and mesh's principal axes with
        matching sign (the naive, unsearched candidate a bare PCA/Umeyama init
        would use) is off by ~166 degrees -- and asserts fit_mesh_to_points,
        which tries all 4 proper-rotation sign candidates and refines each with
        ICP before picking the lowest-residual result, recovers the correct
        pose anyway.
        """
        rng = np.random.default_rng(7)
        mesh = _lopsided_hull_mesh(rng)
        true_rotation = rotvec_to_matrix(rng.uniform(-0.8, 0.8, size=3))
        true_translation = np.array([0.1, -0.05, 0.6])
        view_dir = np.array([1.0, 1.0, 1.0]) / np.sqrt(3.0)
        local_points = _sample_visible_faces(mesh, view_dir, 700, rng)
        points_world = (true_rotation @ local_points.T).T + true_translation
        view_direction_world = true_rotation @ view_dir

        # White-box sanity check: the naive same-sign PCA pairing is indeed
        # badly flipped for this seed, so the test below is actually
        # exercising the multi-start rescue and not passing by accident.
        _, cloud_axes, _ = _pca(points_world)
        _, mesh_axes, _ = _pca(mesh.vertices)
        naive_rotation = cloud_axes @ mesh_axes.T
        assert _rotation_angle_error_deg(true_rotation, naive_rotation) > 90.0

        alignment = fit_mesh_to_points(
            mesh,
            points_world,
            allow_scale=True,
            refine_icp=True,
            rng=np.random.default_rng(107),
            view_direction_world=view_direction_world,
        )
        recovered_rotation = alignment.transform[:3, :3] / alignment.scale
        assert _rotation_angle_error_deg(true_rotation, recovered_rotation) < 5.0

    def test_too_few_points_raises(self) -> None:
        mesh = _box_mesh((0.02, 0.03, 0.04))
        with pytest.raises(AlignmentError):
            fit_mesh_to_points(mesh, np.zeros((3, 3)))

    def test_collinear_points_raise(self) -> None:
        mesh = _box_mesh((0.02, 0.03, 0.04))
        collinear = np.stack([np.linspace(0.0, 1.0, 10)] * 3, axis=1)  # a straight line
        with pytest.raises(AlignmentError):
            fit_mesh_to_points(mesh, collinear)


class TestScaleFromMask:
    """Regression tests pinned to the *measured* demo-object numbers.

    Demo object: a 2x4 LEGO brick, true size ~0.03 m, at 0.63 m from the
    camera. Measured depth MAD inside its (correct, tight) SAM mask was
    0.0265 m -- *larger* than the object -- so the back-projected cloud's own
    spatial extent cannot recover the true size at any filtering strength, but
    the mask's angular extent at the (reliable) median depth can. See the
    module docstrings of :mod:`fpgm.objects.align` and
    :func:`fpgm.objects.align.fit_mesh_to_points` for the full numbers.
    """

    _TRUE_SIZE_M = 0.03
    _MEDIAN_DEPTH_M = 0.63
    _DEPTH_MAD_M = 0.0265
    _FX = _FY = 131.0  # matches the real 320x180 annotation-resolution intrinsics
    _WIDTH, _HEIGHT = 320, 180

    def _noisy_small_object(
        self, rng: np.random.Generator
    ) -> tuple[Camera, np.ndarray, np.ndarray]:
        """A ~3cm object's mask + a depth-noise-dominated point cloud.

        Returns ``(camera, mask, points_world)``. The mask's pixel footprint
        is *derived* from ``_TRUE_SIZE_M``/``_MEDIAN_DEPTH_M``/``_FX`` (not
        hand-picked), so the scenario stays self-consistent if those constants
        are ever tweaked.
        """
        camera = _make_camera(self._WIDTH, self._HEIGHT, fx=self._FX, fy=self._FY)

        bbox_px = max(2, int(round(self._TRUE_SIZE_M * self._FX / self._MEDIAN_DEPTH_M)))
        # Off-centre on purpose: x = (u - cx) * z / fx couples depth noise into
        # X spread proportionally to how far off-axis a pixel is, which is
        # exactly why the *measured* cloud extent (0.215 x 0.185 x 0.121 m)
        # was inflated on every axis, not just Z -- an object dead-centre in
        # frame would only show the effect along Z. Centred at (u - cx) = 0
        # this coupling vanishes and the test would understate the failure
        # mode it exists to document.
        center_u, center_v = self._WIDTH - 20, self._HEIGHT // 2
        x0, y0 = center_u - bbox_px // 2, center_v - bbox_px // 2
        x1, y1 = x0 + bbox_px, y0 + bbox_px

        mask = np.zeros((self._HEIGHT, self._WIDTH), dtype=bool)
        mask[y0:y1, x0:x1] = True
        n = int(mask.sum())

        # Depth noise at the *measured* MAD -- deterministically rescaled
        # (see _rescale_to_mad) so the test doesn't depend on a random draw
        # happening to converge to that MAD at this small sample count.
        noise = _rescale_to_mad(rng.normal(size=n), self._DEPTH_MAD_M)
        depths_mm = np.round((self._MEDIAN_DEPTH_M + noise) * 1000.0).astype(np.uint16)

        depth_mm_image = np.zeros((self._HEIGHT, self._WIDTH), dtype=np.uint16)
        ys, xs = np.nonzero(mask)
        depth_mm_image[ys, xs] = depths_mm

        clip = _DepthClip(initial_depth=depth_mm_image)
        # Default outlier rejection (mad_k, max_depth_span_m) applies, exactly
        # as in the real pipeline -- the measured "+/-1cm trimming still
        # leaves 0.05-0.09 m" finding already accounts for that filtering.
        points_world, _colors = object_point_cloud(
            clip, camera, mask, frame_shape=(self._HEIGHT, self._WIDTH)
        )
        return camera, mask, points_world

    def test_scale_from_mask_recovers_true_size_despite_noise_dominated_cloud(self) -> None:
        rng = np.random.default_rng(42)
        camera, mask, points_world = self._noisy_small_object(rng)

        median_depth = float(np.median(camera.world_to_cam(points_world)[:, 2]))
        mesh = _box_mesh((1.0, 1.0, 1.0))  # unit canonical mesh; scale == the recovered metres

        scale = scale_from_mask(mask, camera, median_depth, mesh)
        assert scale == pytest.approx(self._TRUE_SIZE_M, rel=0.2)

        # The (old, pre-fix) depth-extent estimate: the cloud's own PCA-local
        # extent, matched to the mesh's extent -- exactly what
        # fit_mesh_to_points computed before this fix. Documents *why* the
        # mask path exists: depth noise (MAD 0.0265 m) exceeds the object
        # (0.03 m), so this estimate is wrong by a wide margin.
        _, cloud_axes, _ = _pca(points_world)
        cloud_local = points_world @ cloud_axes
        cloud_extent = cloud_local.max(axis=0) - cloud_local.min(axis=0)
        depth_extent_scale = float(np.mean(cloud_extent / np.asarray(mesh.extent)))
        assert depth_extent_scale > 3.0 * self._TRUE_SIZE_M

    def test_orientation_not_confident_for_noise_dominated_cloud(self) -> None:
        rng = np.random.default_rng(42)
        camera, mask, points_world = self._noisy_small_object(rng)
        mesh = _box_mesh((1.0, 1.0, 1.0))

        alignment = fit_mesh_to_points(
            mesh,
            points_world,
            allow_scale=True,
            refine_icp=True,
            mask_full=mask,
            camera=camera,
            rng=np.random.default_rng(0),
        )

        assert alignment.orientation_confident is False
        assert alignment.scale_source == "mask_angular_extent"
        assert alignment.scale == pytest.approx(self._TRUE_SIZE_M, rel=0.2)
        assert alignment.notes  # a reason was recorded, not a silent fallback
        # "No orientation evidence" must mean literally that: an unrotated
        # placement, not a confident-looking wrong one.
        recovered_rotation = alignment.transform[:3, :3] / alignment.scale
        assert np.allclose(recovered_rotation, np.eye(3), atol=1e-9)

    def test_scale_from_mask_resolution_mismatch_raises(self) -> None:
        camera = _make_camera(320, 180, fx=131.0, fy=131.0)
        wrong_res_mask = np.ones((90, 160), dtype=bool)
        mesh = _box_mesh((1.0, 1.0, 1.0))
        with pytest.raises(ValueError):
            scale_from_mask(wrong_res_mask, camera, 0.63, mesh)


class TestProxyMeshVertexColors:
    """proxy.oriented_box_mesh / convex_hull_mesh should carry through observed colours.

    Every render was a flat grey blob previously (even for a bright teal
    brick) because neither proxy ever set ``ObjectMesh.vertex_colors`` -- see
    the module docstring of :mod:`fpgm.objects.proxy`. Each mesh vertex now
    takes the colour of its *nearest* observed point.
    """

    @staticmethod
    def _colored_cloud(rng: np.random.Generator, n: int = 200) -> tuple[np.ndarray, np.ndarray]:
        points = rng.normal(scale=0.02, size=(n, 3))
        colors = rng.integers(0, 256, size=(n, 3)).astype(np.uint8)
        return points, colors

    def test_oriented_box_mesh_vertex_colors_come_from_observed_points(self) -> None:
        rng = np.random.default_rng(11)
        points, colors = self._colored_cloud(rng)

        mesh = oriented_box_mesh(points, point_colors=colors)

        assert mesh.vertex_colors is not None
        assert mesh.vertex_colors.shape == (mesh.vertices.shape[0], 3)
        assert mesh.vertex_colors.dtype == np.uint8
        observed = {tuple(c) for c in colors}
        assert all(tuple(c) in observed for c in mesh.vertex_colors)

    def test_oriented_box_mesh_without_point_colors_stays_colourless(self) -> None:
        rng = np.random.default_rng(12)
        points, _colors = self._colored_cloud(rng)
        mesh = oriented_box_mesh(points)
        assert mesh.vertex_colors is None

    def test_convex_hull_mesh_vertex_colors_come_from_observed_points(self) -> None:
        rng = np.random.default_rng(13)
        points, colors = self._colored_cloud(rng)

        mesh = convex_hull_mesh(points, point_colors=colors)

        assert mesh.vertex_colors is not None
        assert mesh.vertex_colors.shape == (mesh.vertices.shape[0], 3)
        assert mesh.vertex_colors.dtype == np.uint8
        observed = {tuple(c) for c in colors}
        assert all(tuple(c) in observed for c in mesh.vertex_colors)

    def test_convex_hull_mesh_without_point_colors_stays_colourless(self) -> None:
        rng = np.random.default_rng(14)
        points, _colors = self._colored_cloud(rng)
        mesh = convex_hull_mesh(points)
        assert mesh.vertex_colors is None