orisys-app/apps/studio/tests/test_scene3d_overlays.py

"""Composer overlays and 2D/3D sample-point parity."""

from __future__ import annotations

import numpy as np
from orisys_app_shared.scene3d.composer import SceneComposer
from orisys_app_shared.scene3d.contracts import (
    DisplaySettings,
    FieldDescriptor,
    SurfaceFieldFrame,
    freeze,
)
from orisys_app_shared.scene3d.synthetic import poc_identity, poc_scene
from orisys_app_shared.scene3d.visual_policy import (
    policy_for_settings,
    select_flow_samples,
)


def _depth_frame(assets, *, flow=None, contact_xy=None, fn=0.0):
    h, w = assets.shape_hw
    scalar = freeze(np.zeros((h, w), dtype=np.float32), dtype=np.float32)
    valid = freeze(np.ones((h, w), dtype=bool), dtype=np.bool_)
    flow_uv = None if flow is None else freeze(np.asarray(flow, dtype=np.float32), dtype=np.float32)
    flow_mask = None if flow is None else freeze(np.asarray(valid, dtype=bool).copy(), dtype=np.bool_)
    found = contact_xy is not None
    return SurfaceFieldFrame(
        identity=poc_identity(),
        geometry_id=assets.geometry_id,
        field_domain_id=assets.field_domain_id,
        scalar=scalar,
        valid=valid,
        scalar_descriptor=FieldDescriptor(
            "depth_response", "1", "approximate", False
        ),
        computed_frame_id=0,
        flow_uv=flow_uv,
        flow_valid_mask=flow_mask,
        contact_uv_xy=None if contact_xy is None else tuple(contact_xy),
        contact_found=found,
        contact_valid=found,
        contact_fn=float(fn),
    )


def test_composer_builds_overlays_when_mapping_present() -> None:
    assets = poc_scene()
    h, w = assets.shape_hw
    flow = np.zeros((h, w, 2), dtype=np.float32)
    flow[1, 1] = (0.9, 0.0)
    frame = _depth_frame(assets, flow=flow, contact_xy=(1.0, 1.0), fn=800.0)
    composed = SceneComposer().compose(
        assets,
        frame,
        DisplaySettings(
            flow_arrows_enabled=True,
            contact_arrow_enabled=True,
            flow_grid_spacing=1,
            flow_on_px=0.4,
            flow_off_px=0.3,
            max_arrows=50,
        ),
    )
    assert composed.flow_arrows is not None
    assert composed.flow_arrows.origins.shape[0] >= 1
    assert composed.contact_arrow is not None
    assert composed.contact_arrow.origins.shape[0] == 1


def test_composer_hides_overlays_when_disabled() -> None:
    assets = poc_scene()
    h, w = assets.shape_hw
    flow = np.zeros((h, w, 2), dtype=np.float32)
    flow[1, 1] = (0.9, 0.0)
    frame = _depth_frame(assets, flow=flow, contact_xy=(1.0, 1.0), fn=800.0)
    composed = SceneComposer().compose(
        assets,
        frame,
        DisplaySettings(flow_arrows_enabled=False, contact_arrow_enabled=False),
    )
    assert composed.flow_arrows is None
    assert composed.contact_arrow is None


def test_2d_3d_flow_sample_points_match() -> None:
    """Both figures show the same arrows for the same frame — hysteresis included.

    The regression: the composer kept a "was this sample on last frame" grid and the 2D
    canvas did not, so a sample between ``flow_off_px`` and ``flow_on_px`` was drawn by
    one view and not the other. Measured on a real recording before the fix: 622 arrows
    in the 3D overlay against 432 in the canvas, differing on 34 of 41 frames.

    Frame A is above ``flow_on_px`` (both views draw it); frame B is inside the band —
    above the old hysteresis floor, below the threshold the canvas applies. Both views
    must agree on frame B, which is what fails if the composer re-acquires state the
    canvas does not have.
    """
    assets = poc_scene()
    h, w = assets.shape_hw
    mask = np.ones((h, w), dtype=bool)
    settings = DisplaySettings(
        flow_on_px=0.4,
        flow_off_px=0.3,
        flow_grid_spacing=1,
        max_arrows=50,
    )
    policy = policy_for_settings(settings)
    composer = SceneComposer()

    def both_views(flow_value: float) -> tuple[int, int]:
        flow = np.zeros((h, w, 2), dtype=np.float64)
        flow[1, 1] = (flow_value, 0.0)
        composed = composer.compose(
            assets, _depth_frame(assets, flow=flow), settings
        )
        arrows_3d = (
            0 if composed.flow_arrows is None else int(composed.flow_arrows.origins.shape[0])
        )
        # Exactly what the 2D canvas draws from: ``draw_arrow_display`` calls
        # ``select_flow_samples`` with the policy the window passes it.
        samples_2d = select_flow_samples(flow, mask, policy)
        return arrows_3d, int(samples_2d.rows.size)

    above, above_2d = both_views(0.30)  # |2*flow| = 0.6 px, clearly on
    assert above == 1 and above_2d == 1, (above, above_2d)
    inside_band, inside_band_2d = both_views(0.175)  # |2*flow| = 0.35 px: 0.3 < 0.35 < 0.4
    assert inside_band_2d == 0, inside_band_2d
    assert inside_band == inside_band_2d, (
        "the 3D overlay drew an arrow the 2D canvas did not: the two views are not "
        f"triggering at the same strength ({inside_band} against {inside_band_2d})"
    )


def test_both_views_build_their_policy_from_the_same_settings() -> None:
    """One mapping, one policy: the thresholds are display settings, not per-view ones."""
    settings = DisplaySettings(
        flow_on_px=0.7,
        flow_off_px=0.2,
        flow_grid_spacing=6,
        max_arrows=123,
    )
    policy = policy_for_settings(settings)
    assert (policy.flow_on_px, policy.flow_off_px) == (0.7, 0.2)
    assert policy.flow_grid_spacing == 6
    assert policy.max_arrows == 123
    # Normalised, so a hand-edited preference cannot invert the hysteresis band.
    inverted = policy_for_settings(DisplaySettings(flow_on_px=0.1, flow_off_px=0.9))
    assert inverted.flow_on_px >= inverted.flow_off_px