src/fpgm/depth/scene_flow.py

"""Depth from PointWorld-DROID ``scene_flows`` tracked 3D points, via local interpolation.

There is no downloadable per-pixel depth map for this dataset (the depth archive is
1.23 TB and not cherry-pickable): ``scene_flows`` -- sparse tracked 3D point
positions -- IS the depth source. This module turns those sparse, per-frame 3D
points into dense-enough queryable depth at arbitrary pixel locations via projected
nearest-neighbour interpolation.

Delaunay/barycentric interpolation was considered and rejected: it silently
interpolates *across* occlusion boundaries (e.g. bridging a gripper edge to the
background behind it), producing smoothly-wrong depth with no signal that anything
went wrong. k-NN with an explicit discontinuity-ratio rejection at least produces a
visible ``REJECTED_DISCONTINUITY`` instead of a plausible-looking lie.
"""

from __future__ import annotations

import numpy as np
from scipy.spatial import cKDTree

from fpgm.config import DepthConfig
from fpgm.depth.base import DepthSource
from fpgm.geometry.camera import Camera
from fpgm.types import DepthMethod, DepthQuery, DepthResult, NoValidDepthAnnotationsError

_EPS = 1e-6


def _points_in_mask(uv: np.ndarray, mask: np.ndarray) -> np.ndarray:
    """Return a boolean ``(N,)`` array of whether each ``uv`` falls inside ``mask``.

    Pixel membership **floors**: pixel ``i`` covers the half-open span
    ``[i, i + 1)``. Rounding instead would admit points up to half a pixel *outside*
    the silhouette -- and since this test is what keeps background geometry out of
    the depth interpolation, that half-pixel leak is precisely the contamination the
    mask restriction exists to prevent.
    """
    h, w = mask.shape
    u = np.floor(uv[:, 0]).astype(np.int64)
    v = np.floor(uv[:, 1]).astype(np.int64)
    in_bounds = (u >= 0) & (u < w) & (v >= 0) & (v < h)
    out = np.zeros(uv.shape[0], dtype=bool)
    safe_u = np.where(in_bounds, u, 0)
    safe_v = np.where(in_bounds, v, 0)
    out[in_bounds] = mask[safe_v[in_bounds], safe_u[in_bounds]]
    return out


def _masked_median_rows(values: np.ndarray, mask: np.ndarray) -> np.ndarray:
    """Row-wise median of ``values`` (Q, k), ignoring entries where ``mask`` is False."""
    masked = np.ma.array(values, mask=~mask)
    return np.ma.median(masked, axis=1).filled(np.nan)


def _weighted_median_rows(
    values: np.ndarray, weights: np.ndarray, mask: np.ndarray
) -> np.ndarray:
    """Row-wise weighted median of ``values`` (Q, k) using ``weights``, masked entries excluded.

    A proper weighted median (not a weighted mean): sort each row, walk the
    cumulative weight, and take the value at which cumulative weight first reaches
    half the row's total weight. Robust to a single bad annotation dominating a
    weighted mean.
    """
    w = np.where(mask, weights, 0.0)
    order = np.argsort(values, axis=1)
    sorted_vals = np.take_along_axis(values, order, axis=1)
    sorted_w = np.take_along_axis(w, order, axis=1)
    cum_w = np.cumsum(sorted_w, axis=1)
    total_w = cum_w[:, -1:]
    threshold = total_w / 2.0
    reached = cum_w >= threshold
    first_reached = np.argmax(reached, axis=1)
    return sorted_vals[np.arange(values.shape[0]), first_reached]


class SceneFlowDepthSource(DepthSource):
    """Provides metric depth by projecting ``scene_flows`` and interpolating in image space.

    All HDF5 I/O happens elsewhere; this class only holds already-loaded numpy
    arrays plus the resolved frame convention.
    """

    def __init__(
        self,
        camera: Camera,
        scene_flows: np.ndarray,
        scene_visibility: np.ndarray,
        scene_depth_valid: np.ndarray,
        frame_is_world: bool,
        cfg: DepthConfig,
    ) -> None:
        """Build a scene-flow-backed depth source.

        Args:
            camera: Camera at the scene-flow annotation resolution (rescaled
                internally per query to ``q.query_resolution``).
            scene_flows: ``(T, N, 3)`` tracked 3D point positions.
            scene_visibility: ``(T, N)`` bool, whether each point is visible.
            scene_depth_valid: ``(T, N)`` bool, whether each point's depth annotation
                is trustworthy.
            frame_is_world: Whether ``scene_flows`` positions are WORLD-frame
                (project via ``camera.project``) or CAMERA-frame (``camera.project_cam``).
                Required, not optional -- it must come from
                :func:`fpgm.geometry.convention.detect_scene_flow_convention`, never
                assumed.
            cfg: Interpolation tunables.
        """
        self._camera = camera
        self._scene_flows = np.asarray(scene_flows, dtype=np.float64)
        self._scene_visibility = np.asarray(scene_visibility, dtype=bool)
        self._scene_depth_valid = np.asarray(scene_depth_valid, dtype=bool)
        self._frame_is_world = frame_is_world
        self._cfg = cfg

    def query(self, q: DepthQuery) -> DepthResult:
        """Return interpolated metric depth at ``q.uv`` for frame ``q.frame_idx``.

        Algorithm:
            1. Project that frame's ``scene_flows`` reference points; the support
               set is visible & depth-valid & in-front-of-camera points.
            2. If ``q.object_mask`` is given, restrict support to reference points
               that themselves project inside the mask -- the key defence against a
               background point near in image space but on a different surface
               corrupting the interpolation of a *moving* object's depth. Falls
               back to unrestricted support (with a confidence penalty) if too few
               masked reference points remain.
            3. One ``cKDTree`` built once over the support set, queried in a single
               batched k-NN call -- O((N + Q) log N), never O(Q * N).
            4. Per query point: k-NN within ``max_image_radius_px``; reject for too
               little support, then reject neighbours off the local depth median by
               more than ``depth_discontinuity_ratio`` (occlusion-boundary defence).
            5. Depth = inverse-distance-weighted **median** (not mean) of survivors.

        Raises:
            NoValidDepthAnnotationsError: if the frame has zero usable support
                points at all (distinct from a single query point being rejected).
        """
        cfg = self._cfg
        frame_idx = q.frame_idx
        cam = self._camera.rescaled(*q.query_resolution)

        ref_points = self._scene_flows[frame_idx]
        if self._frame_is_world:
            ref_uv, ref_depth = cam.project(ref_points)
        else:
            ref_uv, ref_depth = cam.project_cam(ref_points)

        base_support = (
            self._scene_visibility[frame_idx]
            & self._scene_depth_valid[frame_idx]
            & (ref_depth > 0)
        )
        if not np.any(base_support):
            raise NoValidDepthAnnotationsError(
                f"frame {frame_idx} has zero scene_flows points with valid "
                "visibility/depth_valid/in-front-of-camera support."
            )

        used_fallback = False
        used_support = base_support
        if q.object_mask is not None:
            in_mask = _points_in_mask(ref_uv, q.object_mask)
            restricted = base_support & in_mask
            if int(np.count_nonzero(restricted)) < cfg.min_support:
                used_fallback = True
                used_support = base_support
            else:
                used_support = restricted

        support_idx = np.flatnonzero(used_support)
        support_uv = ref_uv[support_idx]
        support_depth = ref_depth[support_idx]

        n_query = q.uv.shape[0]
        tree = cKDTree(support_uv)
        k = min(cfg.k_neighbors, support_idx.size)
        dist, nn_idx = tree.query(
            q.uv, k=k, distance_upper_bound=cfg.max_image_radius_px
        )
        if k == 1:
            # cKDTree.query collapses the k-axis for k=1; restore it for uniform
            # downstream shape handling.
            dist = dist.reshape(n_query, 1)
            nn_idx = nn_idx.reshape(n_query, 1)

        valid0 = np.isfinite(dist) & (nn_idx < support_idx.size)
        n_support0 = valid0.sum(axis=1)
        reject_no_support = n_support0 < cfg.min_support

        safe_idx = np.where(valid0, nn_idx, 0)
        depths0 = support_depth[safe_idx]
        median0 = _masked_median_rows(depths0, valid0)
        rel_dev = np.abs(depths0 - median0[:, None]) / np.where(
            median0[:, None] != 0, np.abs(median0[:, None]), _EPS
        )
        valid1 = valid0 & (rel_dev <= cfg.depth_discontinuity_ratio)
        n_support1 = valid1.sum(axis=1)
        reject_discontinuity = (~reject_no_support) & (n_support1 < cfg.min_support)

        ok = ~reject_no_support & ~reject_discontinuity

        weights = 1.0 / (dist + _EPS)
        depth_est = _weighted_median_rows(depths0, weights, valid1)

        mad_rows = np.abs(depths0 - depth_est[:, None])
        mad = _masked_median_rows(mad_rows, valid1)
        normalised_spread = mad / np.where(depth_est != 0, np.abs(depth_est), _EPS)
        confidence = np.clip(
            np.minimum(1.0, n_support1 / max(k, 1)) * (1.0 - normalised_spread), 0.0, 1.0
        )

        out_depth = np.full(n_query, np.nan, dtype=np.float32)
        out_depth[ok] = depth_est[ok].astype(np.float32)
        out_conf = np.zeros(n_query, dtype=np.float32)
        out_conf[ok] = confidence[ok].astype(np.float32)

        method = np.empty(n_query, dtype=object)
        method[:] = DepthMethod.REJECTED_NO_SUPPORT
        method[reject_discontinuity] = DepthMethod.REJECTED_DISCONTINUITY
        method[ok] = (
            DepthMethod.INTERP_UNMASKED_FALLBACK if used_fallback else DepthMethod.INTERP
        )

        if used_fallback:
            out_conf[ok] *= cfg.unmasked_fallback_penalty

        return DepthResult(
            depth=out_depth,
            valid=ok,
            confidence=out_conf,
            method=method,
            n_support=n_support1.astype(np.int32),
        )