src/fpgm/geometry/normals.py

"""Surface normals from a depth map + intrinsics, with no other inputs.

**Design constraint: this must be computable identically from a MuJoCo-rendered
depth buffer at inference time.** The training-time conditioning pipeline
promises that every control channel it feeds a model is something the model
will *also* have available when driven from simulation instead of from a
recorded episode (see the datagen plan's "inference-time contract dictates
training-time conditioning" rule). A learned normal-estimation network, or any
cue derived from the RGB image, would violate that: MuJoCo gives poses and
meshes, not a trained network's opinion of a photograph. So this module depends
on nothing but a depth map and a pinhole intrinsic matrix, and the same two
numpy arrays -- one from a renderer's z-buffer, one from this pipeline's
estimated depth -- produce the same kind of output either way.

Method: unproject the depth map to a camera-frame point map with
:meth:`fpgm.geometry.camera.Camera.unproject_to_cam`, then cross the point
map's central differences. :meth:`Camera.unproject_to_cam` already broadcasts
elementwise over any leading shape (it is called per-*point* elsewhere in this
package, e.g. by :mod:`fpgm.geometry.lifting`, but nothing in its
implementation assumes a 1D point list) -- so no bespoke "grid unprojection"
function was needed; the only new code here is building the pixel-index grid
to feed it and the finite-difference/cross-product step afterwards.
"""

from __future__ import annotations

import numpy as np

from fpgm.geometry.camera import Camera
from fpgm.types import CameraIntrinsics, GeometryError

#: Below this squared length, a cross product is treated as degenerate (the two
#: tangent vectors were parallel or one was zero) rather than normalized, which
#: would otherwise divide by ~0 and produce a numerically huge/garbage unit
#: vector instead of the honest "no normal here".
_MIN_NORMAL_LEN = 1e-12


def _validate_depth_and_valid(
    depth_m: np.ndarray, valid: np.ndarray | None
) -> tuple[np.ndarray, np.ndarray]:
    depth_m = np.asarray(depth_m, dtype=np.float64)
    if depth_m.ndim != 2:
        raise GeometryError(f"depth_m must be (H, W), got shape {depth_m.shape}")
    valid_mask = depth_m > 0
    if valid is not None:
        valid = np.asarray(valid, dtype=bool)
        if valid.shape != depth_m.shape:
            raise GeometryError(
                f"valid must share depth_m's shape {depth_m.shape}, got {valid.shape}"
            )
        valid_mask = valid_mask & valid
    return depth_m, valid_mask


def normals_from_depth(
    depth_m: np.ndarray, K: np.ndarray, *, valid: np.ndarray | None = None
) -> np.ndarray:
    """Unit surface normals in the camera frame, from depth + intrinsics alone.

    Sign convention: the normal points **back toward the camera** for a surface
    facing it. Concretely, for a fronto-parallel plane (constant depth, facing
    the camera along the optical axis), the recovered normal is ``(0, 0, -1)``
    -- this codebase's cameras follow the OpenCV/computer-vision convention
    where depth (``Zc``) increases *away* from the camera along the viewing
    ray (see :meth:`fpgm.geometry.camera.Camera.project_cam`), so "facing the
    camera" is the ``-Z`` direction. This is checked directly in
    ``tests/test_normals.py``.

    Algorithm: unproject every pixel to a camera-frame point using ``K`` alone
    (no extrinsic -- see the module docstring), then at each interior pixel
    take central differences of the point map along the row and column axes and
    cross them: ``normal = normalize(cross(dP/dv, dP/du))``. This is the
    standard "cross of tangent vectors" normal estimator; central differences
    (rather than forward/backward) were chosen because they are the local
    finite-difference stencil with no first-order bias, so a fronto-parallel
    plane recovers exactly its analytic normal rather than one skewed toward
    one neighbour.

    Boundary and invalid handling is deliberately conservative and explicit,
    not silent: a pixel's normal is only computed if the pixel itself *and* all
    four of its immediate neighbours (up/down/left/right) have valid,
    positive depth. Any pixel failing that -- the outer one-pixel image border
    (no neighbour on one side), a hole, or a caller-marked-invalid pixel, or a
    pixel *adjacent* to any of those -- gets the zero vector, never ``NaN``.
    A zero vector cannot be mistaken for a valid unit normal (a real normal
    always has length 1), so "is this normal valid" is simply "is its norm
    nonzero" -- no separate mask needs to be threaded through callers that
    only want the array.

    Args:
        depth_m: ``(H, W)`` float array, metric depth in metres. Values ``<= 0``
            are treated as "no measurement".
        K: ``(3, 3)`` pinhole intrinsic matrix, valid at ``depth_m``'s
            resolution.
        valid: Optional ``(H, W)`` bool array of additional invalidity (e.g. a
            hole-fill provenance mask) to AND with ``depth_m > 0``. If omitted,
            only ``depth_m > 0`` gates validity.

    Returns:
        ``(H, W, 3)`` float32 array of unit normals in the camera frame; the
        zero vector at every pixel described above as invalid.

    Raises:
        GeometryError: if ``depth_m`` is not 2D, ``K`` is not ``(3, 3)``, or
            ``valid`` does not match ``depth_m``'s shape.
    """
    depth_m, valid_mask = _validate_depth_and_valid(depth_m, valid)
    K = np.asarray(K, dtype=np.float64)
    if K.shape != (3, 3):
        raise GeometryError(f"K must be (3, 3), got shape {K.shape}")

    h, w = depth_m.shape
    out = np.zeros((h, w, 3), dtype=np.float32)
    if h < 3 or w < 3:
        # No pixel has a full 4-neighbourhood; every normal is definitionally
        # unavailable rather than an error -- a 1xN or 2xN depth strip is a
        # degenerate but not invalid input.
        return out

    intrinsics = CameraIntrinsics.from_matrix(K, width=w, height=h)
    camera = Camera(intrinsics, np.eye(4))  # identity extrinsic: camera frame only

    v_idx, u_idx = np.meshgrid(np.arange(h, dtype=np.float64), np.arange(w, dtype=np.float64),
                                indexing="ij")
    uv = np.stack([u_idx, v_idx], axis=-1)  # (H, W, 2)
    points = camera.unproject_to_cam(uv, depth_m)  # (H, W, 3), camera frame

    center = valid_mask[1:-1, 1:-1]
    interior_ok = (
        center
        & valid_mask[:-2, 1:-1]  # up (v - 1)
        & valid_mask[2:, 1:-1]  # down (v + 1)
        & valid_mask[1:-1, :-2]  # left (u - 1)
        & valid_mask[1:-1, 2:]  # right (u + 1)
    )

    d_du = (points[1:-1, 2:, :] - points[1:-1, :-2, :]) * 0.5
    d_dv = (points[2:, 1:-1, :] - points[:-2, 1:-1, :]) * 0.5
    normal = np.cross(d_dv, d_du)

    length = np.linalg.norm(normal, axis=-1)
    ok = interior_ok & (length > _MIN_NORMAL_LEN)
    unit = np.zeros_like(normal)
    unit[ok] = normal[ok] / length[ok, None]

    out[1:-1, 1:-1, :] = unit.astype(np.float32)
    return out


def encode_normals_rgb(normals: np.ndarray) -> np.ndarray:
    """Encode unit normals ``[-1, 1]^3`` as the standard 8-bit normal-map RGB.

    ``rgb = round((n * 0.5 + 0.5) * 255)``, i.e. ``-1 -> 0``, ``0 -> 127/128``,
    ``+1 -> 255`` per channel. This is the ubiquitous convention for storing
    normals in an 8-bit image (any renderer/tangent-space normal-map viewer
    assumes it), which is the point: it makes the output paste-compatible with
    every existing viewer and video codec.

    Note that :func:`normals_from_depth`'s zero vector (invalid pixel) encodes
    to mid-grey ``(128, 128, 128)`` here, indistinguishable by colour alone
    from a normal that is genuinely pointing along the camera's X/Y plane with
    zero Z. Callers that need to tell those apart must keep the float array (or
    a separate validity mask) rather than round-tripping through this
    encoding -- see :func:`decode_normals_rgb`.

    Args:
        normals: ``(H, W, 3)`` float array, expected in ``[-1, 1]`` (unit
            vectors, as :func:`normals_from_depth` returns).

    Returns:
        ``(H, W, 3)`` uint8 array.

    Raises:
        GeometryError: if ``normals`` is not ``(H, W, 3)``.
    """
    normals = np.asarray(normals, dtype=np.float32)
    if normals.ndim != 3 or normals.shape[-1] != 3:
        raise GeometryError(f"normals must be (H, W, 3), got shape {normals.shape}")
    rgb = np.clip((normals * 0.5 + 0.5) * 255.0, 0.0, 255.0)
    return np.round(rgb).astype(np.uint8)


def decode_normals_rgb(rgb: np.ndarray) -> np.ndarray:
    """Inverse of :func:`encode_normals_rgb`: 8-bit normal-map RGB -> float ``[-1, 1]``.

    This is a lossy round trip (8 bits per channel quantizes the unit sphere),
    not a re-normalization -- the returned vectors are not guaranteed unit
    length and are not renormalized here, so a caller that needs unit vectors
    back should normalize explicitly after decoding, treating a
    near-mid-grey/near-zero-length result as "was this pixel invalid" per
    :func:`encode_normals_rgb`'s note.

    Args:
        rgb: ``(H, W, 3)`` array (any integer or float dtype) in ``[0, 255]``.

    Returns:
        ``(H, W, 3)`` float32 array in ``[-1, 1]``.

    Raises:
        GeometryError: if ``rgb`` is not ``(H, W, 3)``.
    """
    rgb = np.asarray(rgb)
    if rgb.ndim != 3 or rgb.shape[-1] != 3:
        raise GeometryError(f"rgb must be (H, W, 3), got shape {rgb.shape}")
    return (rgb.astype(np.float32) / 255.0) * 2.0 - 1.0