tests/test_depth.py

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

from __future__ import annotations

import numpy as np
import pytest

from fpgm.config import DepthConfig
from fpgm.depth.scene_flow import SceneFlowDepthSource
from fpgm.geometry.camera import Camera
from fpgm.types import CameraIntrinsics, DepthQuery, NoValidDepthAnnotationsError


def _make_camera() -> Camera:
    intrinsics = CameraIntrinsics(fx=100.0, fy=100.0, cx=50.0, cy=50.0, width=100, height=100)
    return Camera(intrinsics, np.eye(4))  # identity extrinsic: world frame == camera frame


class TestPlanarDepth:
    def test_interpolated_depth_matches_analytic_plane(self):
        camera = _make_camera()
        grid = np.linspace(20.0, 80.0, 15)
        u, v = np.meshgrid(grid, grid)
        ref_uv = np.stack([u.ravel(), v.ravel()], axis=1)
        ref_depth = np.full(ref_uv.shape[0], 2.0)
        ref_pts = camera.unproject_to_cam(ref_uv, ref_depth)  # analytic: constant depth plane

        scene_flows = ref_pts[None, :, :]
        n = ref_pts.shape[0]
        scene_visibility = np.ones((1, n), dtype=bool)
        scene_depth_valid = np.ones((1, n), dtype=bool)

        source = SceneFlowDepthSource(
            camera, scene_flows, scene_visibility, scene_depth_valid, True, DepthConfig()
        )

        rng = np.random.default_rng(3)
        query_uv = rng.uniform(low=25.0, high=75.0, size=(10, 2))
        result = source.query(
            DepthQuery(frame_idx=0, uv=query_uv, query_resolution=(100, 100))
        )

        assert np.all(result.valid)
        assert np.allclose(result.depth, 2.0, atol=1e-3)
        assert np.all(result.confidence > 0.9)  # a perfectly flat plane should be near-certain


class TestObjectMaskRestriction:
    def test_mask_prevents_far_surface_contaminating_near_query(self):
        camera = _make_camera()

        # Foreground "object": a sparse 3x3 grid inside the [45, 55] pixel box, at
        # depth 1.0. Its nearest point to the query below is ~2.24px away.
        object_uv = np.array(
            [[u, v] for u in (45.0, 49.0, 53.0) for v in (45.0, 49.0, 53.0)]
        )
        object_pts = camera.unproject_to_cam(object_uv, np.full(len(object_uv), 1.0))

        # Background: a tight cluster at depth 5.0, positioned just *outside* the
        # object's mask box (u < 45) but closer (~1.1-1.7px) to the query than any
        # object point -- i.e. the far surface annotation intrudes into the k-NN
        # neighbourhood right at the object's silhouette boundary, as it would at a
        # real occlusion edge.
        bg_u = np.linspace(44.3, 44.9, 12)
        bg_v = np.full(12, 47.0)
        bg_uv = np.stack([bg_u, bg_v], axis=1)
        assert np.all(bg_uv[:, 0] < 45.0)  # confirm construction: truly outside the mask box
        bg_pts = camera.unproject_to_cam(bg_uv, np.full(len(bg_uv), 5.0))

        all_pts = np.concatenate([object_pts, bg_pts], axis=0)
        n = all_pts.shape[0]
        scene_flows = all_pts[None, :, :]
        scene_visibility = np.ones((1, n), dtype=bool)
        scene_depth_valid = np.ones((1, n), dtype=bool)

        source = SceneFlowDepthSource(
            camera, scene_flows, scene_visibility, scene_depth_valid, True, DepthConfig()
        )

        # Just inside the object's silhouette, near its boundary -- close enough to
        # the dense background cluster that unmasked k-NN is dominated by it.
        query_uv = np.array([[46.0, 47.0]])
        object_mask = np.zeros((100, 100), dtype=bool)
        object_mask[45:56, 45:56] = True

        unmasked = source.query(
            DepthQuery(frame_idx=0, uv=query_uv, query_resolution=(100, 100))
        )
        masked = source.query(
            DepthQuery(
                frame_idx=0, uv=query_uv, query_resolution=(100, 100), object_mask=object_mask
            )
        )

        assert unmasked.valid[0]
        # All 8 nearest neighbours are the background cluster (closer than any
        # object point), so the unmasked estimate should be ~exactly the far depth.
        assert np.isclose(unmasked.depth[0], 5.0, atol=0.05), (
            "unmasked query should be contaminated by the far surface"
        )

        assert masked.valid[0]
        assert np.isclose(masked.depth[0], 1.0, atol=0.05), "masked query should recover the object"


class TestEmptySupport:
    def test_empty_frame_support_raises(self):
        camera = _make_camera()
        rng = np.random.default_rng(9)
        pts = rng.uniform(low=[-0.3, -0.3, 1.0], high=[0.3, 0.3, 2.0], size=(5, 3))
        scene_flows = pts[None, :, :]
        scene_visibility = np.zeros((1, 5), dtype=bool)  # nothing visible this frame
        scene_depth_valid = np.ones((1, 5), dtype=bool)

        source = SceneFlowDepthSource(
            camera, scene_flows, scene_visibility, scene_depth_valid, True, DepthConfig()
        )

        with pytest.raises(NoValidDepthAnnotationsError):
            source.query(
                DepthQuery(
                    frame_idx=0, uv=np.array([[50.0, 50.0]]), query_resolution=(100, 100)
                )
            )