src/fpgm/geometry/convention.py

"""Detect whether ``scene_flows`` positions are stored in the WORLD or CAMERA frame.

PointWorld-DROID does not document this, and getting it wrong applies a systematic
rigid transform (the extrinsic, or its inverse) to every lifted 3D position while
still producing plausible-looking numbers -- exactly the kind of bug that survives a
casual visual check. So this module never defaults; it scores both hypotheses
against ground truth baked into the same file (in-bounds projection + colour
agreement against the first RGB frame) and refuses to decide when the evidence is
weak.
"""

from __future__ import annotations

import numpy as np

from fpgm.config import ConventionConfig
from fpgm.geometry.camera import Camera
from fpgm.types import (
    AmbiguousConventionError,
    CameraCalibrationMismatchError,
    ConventionDetectionResult,
    FrameConvention,
)


def _bilinear_sample(image: np.ndarray, uv: np.ndarray) -> np.ndarray:
    """Bilinearly sample an ``(H, W, C)`` image at float pixel coordinates ``(N, 2)``.

    Callers must pre-filter ``uv`` to be within ``[0, W-1] x [0, H-1]``; this
    function does not bounds-check.
    """
    h, w = image.shape[:2]
    u = uv[:, 0]
    v = uv[:, 1]
    u0 = np.floor(u).astype(np.int64)
    v0 = np.floor(v).astype(np.int64)
    u1 = np.clip(u0 + 1, 0, w - 1)
    v1 = np.clip(v0 + 1, 0, h - 1)
    u0 = np.clip(u0, 0, w - 1)
    v0 = np.clip(v0, 0, h - 1)
    fu = (u - u0).reshape(-1, 1)
    fv = (v - v0).reshape(-1, 1)

    img = image.astype(np.float64)
    top_left = img[v0, u0]
    top_right = img[v0, u1]
    bottom_left = img[v1, u0]
    bottom_right = img[v1, u1]

    top = top_left * (1 - fu) + top_right * fu
    bottom = bottom_left * (1 - fu) + bottom_right * fu
    return top * (1 - fv) + bottom * fv


def _score_hypothesis(
    uv: np.ndarray,
    depth: np.ndarray,
    colors: np.ndarray,
    initial_rgb: np.ndarray,
    cfg: ConventionConfig,
) -> tuple[float, float, int]:
    """Score one WORLD/CAMERA hypothesis. Returns ``(score, frac_in_bounds, n_scored)``."""
    h, w = initial_rgb.shape[:2]
    in_bounds = (
        (uv[:, 0] >= 0)
        & (uv[:, 0] < w - 1)
        & (uv[:, 1] >= 0)
        & (uv[:, 1] < h - 1)
        & (depth > 0)
    )
    frac_in_bounds = float(np.mean(in_bounds)) if in_bounds.size else 0.0
    n_in_bounds = int(np.count_nonzero(in_bounds))
    if n_in_bounds == 0:
        return 0.0 * cfg.in_bounds_weight, frac_in_bounds, 0
    sampled = _bilinear_sample(initial_rgb, uv[in_bounds])
    reference = colors[in_bounds].astype(np.float64)
    color_agreement = float(np.mean(1.0 - np.abs(sampled - reference) / 255.0))
    score = cfg.in_bounds_weight * frac_in_bounds + cfg.color_weight * color_agreement
    return score, frac_in_bounds, n_in_bounds


def detect_scene_flow_convention(
    camera: Camera,
    scene_flows_frame0: np.ndarray,
    scene_colors_frame0: np.ndarray,
    initial_rgb: np.ndarray,
    cfg: ConventionConfig,
) -> ConventionDetectionResult:
    """Decide whether frame-0 ``scene_flows`` positions are WORLD or CAMERA frame.

    For each hypothesis, points are projected (WORLD via :meth:`Camera.project`,
    CAMERA via :meth:`Camera.project_cam`), the resulting pixel locations are
    checked for being in-bounds and in front of the camera, and the first RGB frame
    is bilinearly sampled at those locations and compared against
    ``scene_colors_frame0``. The hypothesis with the higher combined score wins.

    Note:
        ``initial_rgb`` may be at a different resolution than the annotation
        intrinsic. Callers must pass a ``camera`` already rescaled to
        ``initial_rgb``'s resolution (via :meth:`Camera.rescaled`) -- this function
        does not rescale for you, since silently rescaling would hide a resolution
        bug rather than surface it.

    Args:
        camera: Camera already rescaled to ``initial_rgb``'s resolution.
        scene_flows_frame0: ``(N, 3)`` positions at frame 0, in an unknown frame.
        scene_colors_frame0: ``(N, 3)`` uint8 reference colours at frame 0.
        initial_rgb: ``(H, W, 3)`` uint8 RGB frame, at ``camera.K`` resolution.
        cfg: Scoring weights and decision thresholds.

    Returns:
        The winning :class:`~fpgm.types.ConventionDetectionResult`.

    Raises:
        AmbiguousConventionError: if the two hypotheses' scores are too close to
            call decisively.
        CameraCalibrationMismatchError: if even the winning hypothesis scores badly
            in-bounds -- both hypotheses being bad points at a different failure
            class (wrong intrinsic/extrinsic/serial pairing) that callers should
            triage separately from an ambiguous-but-plausible verdict.
    """
    scene_flows_frame0 = np.asarray(scene_flows_frame0, dtype=np.float64)
    scene_colors_frame0 = np.asarray(scene_colors_frame0)
    n_points = int(scene_flows_frame0.shape[0])

    uv_world, depth_world = camera.project(scene_flows_frame0)
    score_world, frac_world, _n_world = _score_hypothesis(
        uv_world, depth_world, scene_colors_frame0, initial_rgb, cfg
    )

    uv_cam, depth_cam = camera.project_cam(scene_flows_frame0)
    score_cam, frac_cam, _n_cam = _score_hypothesis(
        uv_cam, depth_cam, scene_colors_frame0, initial_rgb, cfg
    )

    margin = abs(score_world - score_cam)
    if score_world >= score_cam:
        winner = FrameConvention.WORLD
        winner_frac = frac_world
    else:
        winner = FrameConvention.CAMERA
        winner_frac = frac_cam

    if margin < cfg.min_margin:
        raise AmbiguousConventionError(
            f"WORLD score={score_world:.4f} vs CAMERA score={score_cam:.4f}: margin "
            f"{margin:.4f} < min_margin {cfg.min_margin}; refusing to guess the "
            "scene_flows frame convention."
        )
    if winner_frac < cfg.min_frac_in_bounds:
        raise CameraCalibrationMismatchError(
            f"winning hypothesis {winner.value} has frac_in_bounds={winner_frac:.4f} "
            f"< min_frac_in_bounds {cfg.min_frac_in_bounds}; both hypotheses project "
            "poorly, suggesting a wrong intrinsic/extrinsic/serial pairing rather "
            "than an ambiguous convention."
        )

    return ConventionDetectionResult(
        convention=winner,
        score_world=score_world,
        score_camera=score_cam,
        margin=margin,
        n_points_scored=n_points,
        frac_in_bounds=winner_frac,
    )