tests/test_depth_fill.py

"""Synthetic, data-free tests for fpgm.depth.fill.

The glass-shard regression lock lives here: :func:`fill_holes_smooth` was
promoted from a script after a nearest-valid-pixel fill (an exact Voronoi
partition of the hole) produced flat polygonal facets that a video diffusion
model rendered as shattered glass. ``TestFillHolesSmooth.test_large_hole_has_no_seams``
constructs that rejected baseline directly and shows it fails a sharp-ridge
check that the promoted push-pull implementation passes, on the same data.
"""

from __future__ import annotations

import cv2
import numpy as np
import pytest
from scipy.ndimage import distance_transform_edt

from fpgm.depth.fill import depth_to_inverse_8bit, fill_holes_smooth, valid_depth
from fpgm.types import GeometryError


def _nearest_fill_baseline(values: np.ndarray, valid: np.ndarray) -> np.ndarray:
    """The fill implementation this module's docstring says not to use.

    Every hole pixel takes the value of its nearest valid pixel -- exactly a
    Voronoi partition of the invalid region, seeded at the valid/invalid
    boundary. Reconstructed here (not imported) because it must not exist
    anywhere in the library; it is only a test fixture proving the regression
    it was replaced for.
    """
    invalid = ~valid
    _, (iy, ix) = distance_transform_edt(invalid, return_indices=True)
    out = values.copy().astype(np.float32)
    out[invalid] = values[iy[invalid], ix[invalid]]
    return out


def _max_interior_laplacian(image: np.ndarray, hole_mask: np.ndarray, erode_px: int = 4) -> float:
    """Peak |Laplacian| strictly inside a hole, away from its boundary with real data.

    Eroding away from the hole boundary excludes the "step" that is expected
    wherever the fill meets real data of a genuinely different value -- the
    thing being measured is whether the *interior* of a filled region is
    smooth (no invented internal seams), not whether the fill matches its
    neighbours at the edge.
    """
    lap = cv2.Laplacian(image.astype(np.float32), cv2.CV_32F, ksize=3)
    kernel = np.ones((2 * erode_px + 1, 2 * erode_px + 1), np.uint8)
    interior = cv2.erode(hole_mask.astype(np.uint8), kernel).astype(bool)
    assert interior.any(), "test construction error: hole too small/thin for the erosion margin"
    return float(np.abs(lap[interior]).max())


def _smooth_plane(h: int = 128, w: int = 128) -> np.ndarray:
    yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
    return 1.0 + 0.01 * xx + 0.02 * yy


class TestFillHolesSmooth:
    def test_valid_pixels_reproduced_verbatim(self):
        rng = np.random.default_rng(0)
        values = rng.uniform(0.2, 4.0, size=(64, 64)).astype(np.float32)
        valid = np.ones((64, 64), dtype=bool)
        valid[20:40, 10:50] = False
        valid[5:12, 5:12] = False

        out = fill_holes_smooth(values, valid)

        assert np.array_equal(out[valid], values[valid]), (
            "measured pixels must survive the fill bit-for-bit, never blended with "
            "the interpolant"
        )

    def test_no_holes_returns_input_values(self):
        values = np.full((16, 16), 3.5, dtype=np.float32)
        valid = np.ones((16, 16), dtype=bool)
        out = fill_holes_smooth(values, valid)
        assert np.array_equal(out, values)

    def test_large_hole_has_no_seams_unlike_nearest_fill_baseline(self):
        """The glass-shard regression lock.

        A large contiguous hole (~35% of the image, split across two blobs, like
        the real capture that triggered this) is filled two ways from the same
        smooth analytic plane: the promoted push-pull implementation, and the
        nearest-valid-pixel baseline that was rejected for production use.
        Only the baseline is expected to contain a sharp internal ridge --
        that ridge *is* the Voronoi seam that rendered as a glass shard.
        """
        values = _smooth_plane()
        valid = np.ones(values.shape, dtype=bool)
        valid[30:90, 20:100] = False
        valid[10:25, 60:110] = False
        hole = ~valid
        assert hole.mean() > 0.30  # matches "large contiguous hole" from the bug report

        filled_pushpull = fill_holes_smooth(values, valid)
        filled_nearest = _nearest_fill_baseline(values, valid)

        # Both fills must still reproduce the measured pixels exactly -- a
        # sanity check that the baseline itself is a faithful reconstruction of
        # the rejected algorithm, not an accidentally-different one.
        assert np.array_equal(filled_pushpull[valid], values[valid])
        assert np.array_equal(filled_nearest[valid], values[valid])

        pushpull_ridge = _max_interior_laplacian(filled_pushpull, hole)
        nearest_ridge = _max_interior_laplacian(filled_nearest, hole)

        # The rejected algorithm must actually exhibit the seam being guarded
        # against, on this exact data -- otherwise this test would pass for the
        # wrong reason.
        assert nearest_ridge > 1.0, (
            f"nearest-fill baseline did not produce a sharp seam ({nearest_ridge=}); "
            "the test data no longer reproduces the failure mode being regression-locked"
        )
        # The promoted implementation must be smooth in the same region.
        assert pushpull_ridge < 0.5
        # And decisively smoother than the baseline it replaced -- not just
        # under some absolute threshold, but by a wide margin on the identical
        # input.
        assert pushpull_ridge < nearest_ridge / 10.0

    def test_all_invalid_raises(self):
        values = np.ones((8, 8), dtype=np.float32)
        valid = np.zeros((8, 8), dtype=bool)
        with pytest.raises(GeometryError):
            fill_holes_smooth(values, valid)

    def test_shape_mismatch_raises(self):
        values = np.ones((8, 8), dtype=np.float32)
        valid = np.ones((8, 9), dtype=bool)
        with pytest.raises(GeometryError):
            fill_holes_smooth(values, valid)


class TestValidDepth:
    def test_zero_depth_is_invalid(self):
        depth = np.array([[0.0, 1.0], [2.0, 0.0]], dtype=np.float32)
        seg = np.zeros((2, 2), dtype=np.uint8)
        out = valid_depth(depth, seg, min_background_m=0.0)
        assert np.array_equal(out, [[False, True], [True, False]])

    def test_only_background_rejected_by_min_depth_floor(self):
        # A 2x2 grid: background (seg=0) near/far, foreground (seg=1) near/far.
        depth = np.array([[0.10, 5.0], [0.10, 5.0]], dtype=np.float32)
        seg = np.array([[0, 0], [1, 1]], dtype=np.uint8)  # top row bg, bottom row fg

        out = valid_depth(depth, seg, min_background_m=0.40)

        # Background near reading (0.10 m < 0.40 m floor) is rejected.
        assert not out[0, 0]
        assert out[0, 1]
        # Foreground is exempt from the floor even though it reads the same
        # 0.10 m as the rejected background pixel.
        assert out[1, 0]
        assert out[1, 1]

    def test_min_depth_zero_disables_background_floor(self):
        depth = np.array([[0.01, 5.0]], dtype=np.float32)
        seg = np.array([[0, 0]], dtype=np.uint8)
        out = valid_depth(depth, seg, min_background_m=0.0)
        assert np.array_equal(out, [[True, True]])

    def test_shape_mismatch_raises(self):
        depth = np.zeros((4, 4), dtype=np.float32)
        seg = np.zeros((4, 5), dtype=np.uint8)
        with pytest.raises(GeometryError):
            valid_depth(depth, seg, min_background_m=0.0)


class TestDepthToInverse8Bit:
    def test_monotonic_nearer_is_brighter(self):
        depths = np.array([0.5, 1.0, 2.0, 4.0, 8.0], dtype=np.float32)  # far -> near order below
        depths_far_to_near = depths[::-1]  # 8, 4, 2, 1, 0.5 metres: nearer each step
        inv = 1.0 / depths_far_to_near
        lo, hi = float(inv.min()), float(inv.max())

        out = depth_to_inverse_8bit(depths_far_to_near, lo, hi)

        assert np.all(np.diff(out.astype(np.int32)) > 0), (
            "brightness must increase as depth decreases"
        )

    def test_endpoint_values(self):
        lo, hi = 0.1, 1.0
        far_depth = np.array([[1.0 / lo]], dtype=np.float32)
        near_depth = np.array([[1.0 / hi]], dtype=np.float32)

        assert depth_to_inverse_8bit(far_depth, lo, hi)[0, 0] == 0
        assert depth_to_inverse_8bit(near_depth, lo, hi)[0, 0] == 255

    def test_range_clamping(self):
        lo, hi = 0.5, 1.0
        beyond_far = np.array([[1.0 / 0.1]], dtype=np.float32)  # inverse depth < lo
        beyond_near = np.array([[1.0 / 10.0]], dtype=np.float32)  # inverse depth > hi

        # beyond_far metres -> huge depth -> tiny inverse -> below lo -> clamps to 0
        assert depth_to_inverse_8bit(beyond_far, lo, hi)[0, 0] == 0
        # beyond_near metres -> tiny depth -> huge inverse -> above hi -> clamps to 255
        assert depth_to_inverse_8bit(beyond_near, lo, hi)[0, 0] == 255

    def test_dtype_and_shape_preserved(self):
        depth = np.full((3, 5), 2.0, dtype=np.float32)
        out = depth_to_inverse_8bit(depth, 0.1, 2.0)
        assert out.dtype == np.uint8
        assert out.shape == (3, 5)

    def test_degenerate_range_raises(self):
        depth = np.ones((2, 2), dtype=np.float32)
        with pytest.raises(GeometryError):
            depth_to_inverse_8bit(depth, 1.0, 1.0)
        with pytest.raises(GeometryError):
            depth_to_inverse_8bit(depth, 1.0, 0.5)