src/fpgm/depth/fill.py

"""Depth hole-filling and 8-bit encoding, promoted from ``scripts/export_cosmos_input.py``.

These three functions were first written as private helpers in the Cosmos export
script and were battle-tested there against a real production failure (see
:func:`fill_holes_smooth`). They are promoted here verbatim-in-behaviour so the
next consumer -- the datagen pipeline's dense-depth stage -- gets the fix for
free instead of re-discovering it.
"""

from __future__ import annotations

import cv2
import numpy as np

from fpgm.types import GeometryError


def fill_holes_smooth(values: np.ndarray, valid: np.ndarray, levels: int = 7) -> np.ndarray:
    """Fill 2D holes by Gaussian-pyramid push-pull, keeping measured pixels exact.

    **Do not replace this with a nearest-valid-pixel fill.** That was the first
    implementation and it produced a visible, serious artifact in a generated
    video: assigning every hole pixel the value of its nearest valid pixel is
    *precisely* a Voronoi partition of the image plane, so a filled region breaks
    into flat polygonal cells meeting along straight seams. With ~17% of pixels
    holed and the holes forming large contiguous blobs, that painted a
    tessellation of fake planar facets into the depth control channel -- and a
    video diffusion model conditioned on it faithfully rendered those facets as
    shattered-glass surfaces hanging in the scene. It is nearly invisible in a
    downscaled preview of the depth video and obvious the moment an edge
    detector (or a Laplacian) is run over it, which is how it was finally caught.

    Push-pull instead: decimate value and coverage through a Gaussian pyramid,
    then reconstruct coarse-to-fine, letting each level supply only what the
    finer level lacks. The interpolant is smooth by construction, so it adds no
    edges that were not already in the data -- there is no length scale at which
    it can produce a straight seam. Measured pixels are written back verbatim
    afterwards: the fill only ever invents values where there were none, and it
    invents them smoothly, which is the honest representation of "we don't know"
    for a depth (or any other per-pixel scalar) surface.

    Args:
        values: ``(H, W)`` float array of measurements. Values at invalid
            locations are ignored (never read) -- they do not need to be
            pre-zeroed by the caller.
        valid: ``(H, W)`` bool array, ``True`` where ``values`` is a real
            measurement.
        levels: Depth of the Gaussian pyramid. 7 halves the resolution 7 times,
            which for typical camera resolutions (VGA to 1080p) reaches a coarse
            level small enough that even a large contiguous hole has *some*
            coverage to pull from.

    Returns:
        ``(H, W)`` float32 array, equal to ``values`` at ``valid`` pixels and a
        smooth fill elsewhere.

    Raises:
        GeometryError: if ``values`` and ``valid`` disagree in shape, if either
            is not 2D, or if ``valid`` is all-``False`` (there is nothing to seed
            the fill from, so any output would be pure invention rather than an
            interpolation of real data).
    """
    values = np.asarray(values)
    valid = np.asarray(valid, dtype=bool)
    if values.ndim != 2:
        raise GeometryError(f"values must be (H, W), got shape {values.shape}")
    if values.shape != valid.shape:
        raise GeometryError(
            f"values and valid must share shape, got {values.shape} vs {valid.shape}"
        )
    if not valid.any():
        raise GeometryError(
            "valid is all-False: there are no measured pixels to fill from, so any "
            "output here would be pure invention rather than an interpolation."
        )

    if valid.all():
        return values.astype(np.float32)

    m = valid.astype(np.float32)
    vs, ms = [values.astype(np.float32) * m], [m]
    for _ in range(levels):
        vs.append(cv2.pyrDown(vs[-1]))
        ms.append(cv2.pyrDown(ms[-1]))

    up = vs[-1] / np.maximum(ms[-1], 1e-6)
    for k in range(len(vs) - 2, -1, -1):
        up = cv2.resize(up, (vs[k].shape[1], vs[k].shape[0]), interpolation=cv2.INTER_LINEAR)
        w = np.clip(ms[k], 0.0, 1.0)
        up = vs[k] / np.maximum(ms[k], 1e-6) * w + up * (1.0 - w)

    out = up
    out[valid] = values[valid]
    return out.astype(np.float32)


def valid_depth(depth_m: np.ndarray, seg: np.ndarray, min_background_m: float) -> np.ndarray:
    """Which depth pixels are trustworthy measurements.

    Zero means "no data" by the renderer's convention. On top of that, a
    *background* reading nearer than ``min_background_m`` is rejected: on the ZED
    stereo capture this was promoted from, the sensor's minimum reported depth is
    233 mm, and stereo matching failure saturates disparity -- so it lands at the
    sensor's *near* limit rather than anywhere near the truth. A cluster of
    background readings pinned just above that minimum, concentrated in dark or
    textureless regions (a black curtain, a far wall), is a failure mode, not a
    measurement: left in, those pixels become the *brightest* thing in an
    inverse-depth control (see :func:`depth_to_inverse_8bit`), drawing a hard
    bogus structure across the frame and dragging the normalisation range so
    everything real loses contrast.

    The foreground is exempt on purpose: a robot arm or manipulated object
    genuinely can come within centimetres of a wide-angle camera, and its depth
    typically comes from a mesh renderer's z-buffer, which this stereo-failure
    filter has no business touching.

    Args:
        depth_m: ``(H, W)`` float array, metres.
        seg: ``(H, W)`` integer array of segmentation ids, ``0`` = background.
        min_background_m: Background depth readings nearer than this are treated
            as stereo-matching failures rather than measurements. ``0`` disables
            the background-specific rejection (only ``depth_m > 0`` still
            applies).

    Returns:
        ``(H, W)`` bool array, ``True`` where ``depth_m`` should be trusted.

    Raises:
        GeometryError: if ``depth_m`` and ``seg`` disagree in shape.
    """
    depth_m = np.asarray(depth_m)
    seg = np.asarray(seg)
    if depth_m.shape != seg.shape:
        raise GeometryError(
            f"depth_m and seg must share shape, got {depth_m.shape} vs {seg.shape}"
        )

    valid = depth_m > 0
    if min_background_m > 0:
        valid &= ~((seg == 0) & (depth_m < min_background_m))
    return valid


def depth_to_inverse_8bit(depth_m: np.ndarray, lo: float, hi: float) -> np.ndarray:
    """Metric metres -> 8-bit relative inverse depth (near = bright).

    This matches the convention used by relative-depth estimators such as
    DepthAnything, which is what Cosmos-Transfer2.5's own depth control channel
    was trained against -- so a control video built this way looks, to the model,
    like the output of its native depth extractor rather than an unfamiliar
    signal.

    **Normalisation must be computed once, globally, for the whole clip -- never
    per frame.** ``lo``/``hi`` are the *clip-wide* robust percentiles of inverse
    depth, computed by the caller once and passed in unchanged for every frame.
    Recomputing them per frame is the obvious-looking implementation and it is
    wrong: as the scene's near/far extremes drift frame to frame, a per-frame
    range makes the encoded video's brightness drift with it, and a video
    diffusion model conditioned on that control faithfully reproduces the drift
    as visible brightness pumping in the generated output. A single range fixed
    across the clip costs nothing to compute and removes the failure mode
    entirely. Use robust percentiles (not min/max) when deriving ``lo``/``hi``,
    so a single stray near/far pixel cannot crush the usable range for every
    other frame in the clip.

    Args:
        depth_m: Metric depth, any shape, metres. Non-positive values are
            clamped away from a division singularity (``1e-6`` metres) rather
            than propagating ``inf``/``NaN`` -- callers are expected to have
            already filled or masked non-measurements (see
            :func:`fill_holes_smooth`, :func:`valid_depth`).
        lo: Clip-wide low percentile of inverse depth (1/metres), i.e. the *far*
            end of the range.
        hi: Clip-wide high percentile of inverse depth (1/metres), i.e. the
            *near* end of the range. Must be strictly greater than ``lo``.

    Returns:
        uint8 array, same shape as ``depth_m``, in ``[0, 255]``; ``0`` at
        ``lo`` (far), ``255`` at ``hi`` (near), clamped outside that range.

    Raises:
        GeometryError: if ``hi <= lo``, since that range is degenerate -- every
            pixel would silently collapse to the same output value instead of
            visibly failing.
    """
    if not (hi > lo):
        raise GeometryError(f"depth_to_inverse_8bit requires hi > lo, got lo={lo} hi={hi}")

    depth_m = np.asarray(depth_m, dtype=np.float32)
    inv = 1.0 / np.maximum(depth_m, 1e-6)
    norm = np.clip((inv - lo) / (hi - lo), 0.0, 1.0)
    return (norm * 255.0).astype(np.uint8)