tests/test_normals.py

"""Synthetic, data-free tests for fpgm.geometry.normals."""

from __future__ import annotations

import numpy as np
import pytest

from fpgm.geometry.normals import decode_normals_rgb, encode_normals_rgb, normals_from_depth
from fpgm.types import GeometryError

_FX = _FY = 100.0
_CX = _CY = 50.0
_H = _W = 101


def _K() -> np.ndarray:
    return np.array([[_FX, 0.0, _CX], [0.0, _FY, _CY], [0.0, 0.0, 1.0]])


class TestFrontoParallelPlane:
    def test_normal_is_minus_z(self):
        """Sign convention: a plane facing the camera recovers normal (0, 0, -1).

        Chosen because this codebase's cameras follow the OpenCV convention
        where depth increases away from the camera along +Z (see
        ``Camera.project_cam``), so the direction pointing back toward the
        camera -- i.e. "facing the viewer" -- is -Z.
        """
        depth = np.full((_H, _W), 2.0, dtype=np.float32)
        normals = normals_from_depth(depth, _K())

        interior = normals[10:-10, 10:-10]
        expected = np.array([0.0, 0.0, -1.0], dtype=np.float32)
        assert np.allclose(interior, expected, atol=1e-6)


class TestTiltedPlane:
    def test_recovered_normal_matches_analytic_plane(self):
        """A known-tilted plane's depth map is generated analytically and inverted.

        The plane is defined by an implicit equation ``coef . P = d0`` in the
        camera frame, with ``coef`` chosen as ``-n_expected`` so that
        ``n_expected`` is exactly what :func:`TestFrontoParallelPlane`'s sign
        convention predicts should come back out (the fronto-parallel case is
        the ``coef = (0, 0, 1)``, ``n_expected = (0, 0, -1)`` special case of
        this same construction).
        """
        n_expected = np.array([0.3, 0.2, -0.9])
        n_expected /= np.linalg.norm(n_expected)
        coef = -n_expected
        d0 = 2.0

        vv, uu = np.mgrid[0:_H, 0:_W].astype(np.float64)
        denom = coef[0] * (uu - _CX) / _FX + coef[1] * (vv - _CY) / _FY + coef[2]
        depth = d0 / denom
        assert np.all(depth > 0), "test construction error: plane must stay in front of the camera"

        normals = normals_from_depth(depth, _K())

        interior = normals[10:-10, 10:-10]
        err = np.linalg.norm(interior - n_expected.astype(np.float32), axis=-1)
        assert err.max() < 1e-3


class TestInvalidHandling:
    def test_zero_depth_pixel_and_neighbours_are_zero_not_nan(self):
        depth = np.full((_H, _W), 2.0, dtype=np.float32)
        depth[50, 50] = 0.0  # a single dropped pixel in an otherwise flat plane

        normals = normals_from_depth(depth, _K())

        assert not np.isnan(normals).any()
        # The hole pixel itself, and every pixel that depends on it as an
        # immediate neighbour in the central-difference stencil, must be the
        # zero vector.
        for (i, j) in [(50, 50), (49, 50), (51, 50), (50, 49), (50, 51)]:
            assert np.array_equal(normals[i, j], [0.0, 0.0, 0.0])
        # A pixel two steps away is unaffected.
        assert np.allclose(normals[47, 50], [0.0, 0.0, -1.0], atol=1e-6)

    def test_explicit_valid_mask_is_honoured(self):
        depth = np.full((20, 20), 2.0, dtype=np.float32)
        valid = np.ones((20, 20), dtype=bool)
        valid[10, 10] = False  # depth > 0 here, but explicitly marked invalid

        normals = normals_from_depth(depth, _K(), valid=valid)
        assert np.array_equal(normals[10, 10], [0.0, 0.0, 0.0])

    def test_border_pixels_are_zero(self):
        depth = np.full((20, 20), 2.0, dtype=np.float32)
        normals = normals_from_depth(depth, _K())
        assert np.array_equal(normals[0, :], np.zeros((20, 3)))
        assert np.array_equal(normals[-1, :], np.zeros((20, 3)))
        assert np.array_equal(normals[:, 0], np.zeros((20, 3)))
        assert np.array_equal(normals[:, -1], np.zeros((20, 3)))

    def test_too_small_input_returns_all_zero(self):
        depth = np.full((2, 2), 1.0, dtype=np.float32)
        normals = normals_from_depth(depth, _K())
        assert normals.shape == (2, 2, 3)
        assert np.array_equal(normals, np.zeros((2, 2, 3)))


class TestValidation:
    def test_bad_depth_ndim_raises(self):
        with pytest.raises(GeometryError):
            normals_from_depth(np.zeros((4, 4, 4)), _K())

    def test_bad_k_shape_raises(self):
        with pytest.raises(GeometryError):
            normals_from_depth(np.zeros((4, 4)), np.eye(4))

    def test_valid_shape_mismatch_raises(self):
        with pytest.raises(GeometryError):
            normals_from_depth(np.zeros((4, 4)), _K(), valid=np.ones((5, 5), dtype=bool))


class TestEncodeDecodeRoundTrip:
    def test_encode_known_values(self):
        normals = np.array([[[0.0, 0.0, -1.0], [1.0, -1.0, 0.0]]], dtype=np.float32)
        rgb = encode_normals_rgb(normals)
        assert rgb.dtype == np.uint8
        assert tuple(rgb[0, 0]) == (128, 128, 0)
        assert tuple(rgb[0, 1]) == (255, 0, 128)

    def test_decode_inverts_encode_within_quantization(self):
        rng = np.random.default_rng(1)
        raw = rng.normal(size=(8, 8, 3))
        normals = raw / np.linalg.norm(raw, axis=-1, keepdims=True)

        rgb = encode_normals_rgb(normals)
        decoded = decode_normals_rgb(rgb)

        # 8-bit quantization: within one LSB (2/255) per channel.
        assert np.abs(decoded - normals).max() < 2.0 / 255.0 + 1e-6

    def test_encode_bad_shape_raises(self):
        with pytest.raises(GeometryError):
            encode_normals_rgb(np.zeros((4, 4, 2)))

    def test_decode_bad_shape_raises(self):
        with pytest.raises(GeometryError):
            decode_normals_rgb(np.zeros((4, 4)))