NODEDC_MISSION_CORE/tests/test_geometry_semantic_roi.py

from __future__ import annotations

from pathlib import Path

import numpy as np
from PIL import Image

from k1link.perception.contracts import (
    EvidenceBasis,
    EvidenceCurrentness,
    MetricGeometry,
    ObstacleObservation,
)
from k1link.perception.geometry import GeometryFrame
from k1link.perception.geometry_math import Kb4ProjectionProfile
from k1link.perception.geometry_semantic_roi import (
    build_geometry_semantic_rois,
    materialize_geometry_semantic_crop,
    select_geometry_roi_detections,
)
from k1link.perception.object_understanding import load_object_semantic_vocabulary
from k1link.perception.open_vocabulary_semantics import (
    OpenVocabularyDetection,
    fuse_open_vocabulary_detections,
    load_open_vocabulary_semantic_profile,
)

REPOSITORY_ROOT = Path(__file__).resolve().parents[1]
PROFILE_PATH = REPOSITORY_ROOT / "config/perception/open-vocabulary-semantic-shadow-v0.json"
VOCABULARY_PATH = REPOSITORY_ROOT / "config/perception/object-semantic-vocabulary-v0.json"


def _point_for_pixel(u: float, v: float, *, z: float = 5.0) -> tuple[float, float, float]:
    theta_x = (u - 400.0) / 100.0
    theta_y = (v - 300.0) / 100.0
    ray = np.asarray((np.tan(theta_x), np.tan(theta_y), 1.0), dtype=np.float64)
    return tuple(float(item) for item in ray * z)  # type: ignore[return-value]


def _frame() -> GeometryFrame:
    points = np.asarray(
        (
            _point_for_pixel(350.0, 250.0),
            _point_for_pixel(450.0, 250.0),
            _point_for_pixel(350.0, 350.0),
            _point_for_pixel(450.0, 350.0),
            _point_for_pixel(600.0, 300.0),
        ),
        dtype=np.float64,
    )
    return GeometryFrame(
        frame_index=121,
        points_map=points,
        point_class=np.full(points.shape[0], 2, dtype=np.uint8),
        sensor_position_map=np.zeros(3, dtype=np.float64),
        sensor_orientation_xyzw=np.asarray((0.0, 0.0, 0.0, 1.0), dtype=np.float64),
        projection=Kb4ProjectionProfile(
            width=800,
            height=600,
            intrinsic_fx_fy_cx_cy=(100.0, 100.0, 400.0, 300.0),
            distortion_kb4=(0.0, 0.0, 0.0, 0.0),
            t_camera_from_lidar=np.eye(4, dtype=np.float64),
        ),
        surface_valid=True,
    )


def _observation(ordinal: int, source_point_ids: tuple[int, ...]) -> ObstacleObservation:
    return ObstacleObservation(
        observation_id=f"frame-000121:observation-{ordinal}",
        occupancy_key=f"frame-000121:occupancy-{ordinal}",
        source_id="RAVNOVES00",
        frame_id="frame-000121",
        evidence_time_ns=121,
        basis=EvidenceBasis.LIDAR,
        currentness=EvidenceCurrentness.CURRENT,
        occupied_support=True,
        source_point_ids=source_point_ids,
        metric_geometry=MetricGeometry(
            coordinate_frame="map",
            centroid_xyz_m=(1.0, 2.0, 0.5),
            range_m=2.2,
            covariance_diagonal_m2=(0.1, 0.1, 0.1),
        ),
        proposal_ids=(),
        semantic_hint=None,
        reason_codes=("test-geometry-only",),
    )


def test_geometry_points_own_crop_and_sparse_observation_stays_unprojected(
    tmp_path: Path,
) -> None:
    result = build_geometry_semantic_rois(
        frame=_frame(),
        observations=(_observation(0, (0, 1, 2, 3)), _observation(1, (4,))),
    )

    assert len(result.rois) == 1
    assert result.not_projected_observations[0].observation_id.endswith("observation-1")
    roi = result.rois[0]
    assert roi.observation.source_point_ids == (0, 1, 2, 3)
    assert roi.core_region.x_min < 400.0 < roi.core_region.x_max
    assert roi.core_region.y_min < 300.0 < roi.core_region.y_max
    assert roi.crop_region.x_min <= roi.core_region.x_min
    assert roi.crop_region.y_max >= roi.core_region.y_max

    source = tmp_path / "frame-000121.png"
    Image.new("RGB", (800, 600), (114, 114, 114)).save(source)
    destination = tmp_path / roi.crop_name
    materialize_geometry_semantic_crop(
        image_path=source,
        roi=roi,
        destination=destination,
    )
    with Image.open(destination) as crop:
        assert crop.size == (
            int(roi.crop_region.x_max - roi.crop_region.x_min),
            int(roi.crop_region.y_max - roi.crop_region.y_min),
        )


def test_roi_selection_returns_only_cluster_covering_geometry_core() -> None:
    profile = load_open_vocabulary_semantic_profile(PROFILE_PATH)
    vocabulary = load_object_semantic_vocabulary(VOCABULARY_PATH)
    roi = build_geometry_semantic_rois(
        frame=_frame(),
        observations=(_observation(0, (0, 1, 2, 3)),),
    ).rois[0]
    fusion = fuse_open_vocabulary_detections(
        (
            OpenVocabularyDetection(
                detection_id="inside",
                source_id="RAVNOVES00",
                frame_id="frame-000121",
                prompt_set_id="urban-static/v0",
                raw_label="trash bin",
                confidence=0.8,
                region=roi.core_region,
            ),
            OpenVocabularyDetection(
                detection_id="outside",
                source_id="RAVNOVES00",
                frame_id="frame-000121",
                prompt_set_id="urban-static/v0",
                raw_label="traffic cone",
                confidence=0.9,
                region=type(roi.core_region)(10.0, 10.0, 50.0, 50.0),
            ),
        ),
        profile=profile,
        vocabulary=vocabulary,
        valid_fov_mask=np.ones((600, 800), dtype=np.bool_),
    )

    selected = select_geometry_roi_detections(roi, fusion)

    assert tuple(item.detection_id for item in selected) == ("inside",)