scripts/verify_world_frame.py

#!/usr/bin/env python
"""Verify PointWorld-DROID ``scene_flows`` from two cameras share one world frame.

This is the single most load-bearing assumption of the datagen pipeline planned
on top of this dataset: that camera ext1's and camera ext2's ``scene_flows`` for
the *same clip* can be normalized (via :func:`fpgm.geometry.convention.detect_scene_flow_convention`
+ each clip's own ``extrinsic``) into one common world frame, so a single
canonical point cloud/robot-base frame can be built from either camera
indifferently. Getting this wrong would mean every downstream stage (dense
depth, object discovery, pose fitting) silently mixes two different frames.

Four independent checks, each with its own PASS/FAIL verdict:

1. Convention detection succeeds decisively for both cameras (reuses the
   existing, already-adversarially-designed detector -- this script does not
   re-implement frame-guessing).
2. Cross-camera colour agreement: frame-0 world points from camera A, projected
   into camera B, should sample colours in B's ``initial_rgb`` close to what A
   itself recorded (``scene_colors``) -- and vice versa. This is exactly the
   convention detector's own scoring idea, just applied *across* cameras
   instead of within one.
3. Cross-camera depth agreement: the same reprojected points' camera-B depth
   should agree with B's own ``initial_depth`` at that pixel.
4. The rigid transform between the two cameras implied by their ``extrinsic``
   matrices (as read from the clip) should match the transform implied by
   ``optimized_extrinsics`` in ``<uuid>_cameras.json`` -- an independent
   calibration source for the same physical rig.

If cameras do NOT agree, this prints a loud FAIL with the actual numbers rather
than silently downgrading -- the whole one-world-frame design in the datagen
plan depends on this holding.

Usage:
    PYTHONPATH=src python scripts/verify_world_frame.py \\
        --episode AUTOLab+0d4edc83+2023-10-21-19h-07m-04s --clip 20:31 \\
        --camera-a 22008760 --camera-b 24400334
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

import numpy as np  # noqa: E402

from fpgm.config import ConventionConfig  # noqa: E402
from fpgm.data.pointworld import FlowsReader, PointWorldStore  # noqa: E402
from fpgm.geometry.camera import Camera  # noqa: E402
from fpgm.geometry.convention import detect_scene_flow_convention  # noqa: E402
from fpgm.geometry.transforms import invert_se3, matrix_to_rotvec  # noqa: E402
from fpgm.types import (  # noqa: E402
    AmbiguousConventionError,
    CameraCalibrationMismatchError,
    DataError,
    FrameConvention,
    SceneFlowClip,
)
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

logger = get_logger("verify_world_frame")

# Thresholds chosen to separate "same world frame, ordinary sensor/JPEG/lighting
# noise" from "off by a rigid transform" (which is not a subtle effect -- a wrong
# frame typically blows these numbers up by 5-10x, not by a few percent). They are
# deliberately loose, not scientific: this script's job is to catch a categorical
# frame bug, not to certify sub-pixel calibration quality.
_PASS_COLOR_DIFF_PER_CHANNEL = 30.0  # 0-255 scale
_PASS_DEPTH_ERROR_MM = 80.0
_PASS_TRANSLATION_ERROR_M = 0.01
_PASS_ROTATION_ERROR_DEG = 1.0
#: DROID's two exterior cameras only partially overlap (they are placed to see
#: the scene from different angles, not to stereo-overlap fully), so most of one
#: camera's points legitimately fall outside the other's frustum. A small but
#: nonzero fraction landing in-bounds is the expected signature of "same frame,
#: partial overlap"; near-zero would instead suggest a frame/transform bug.
_PASS_MIN_FRAC_IN_BOUNDS = 0.02


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--episode", required=True, help="DROID episode uuid")
    parser.add_argument("--clip", required=True, help='PointWorld clip key, e.g. "20:31"')
    parser.add_argument("--camera-a", required=True, help="camera serial, e.g. 22008760")
    parser.add_argument("--camera-b", required=True, help="camera serial, e.g. 24400334")
    parser.add_argument(
        "--data-dir", type=Path, default=REPO_ROOT / "data" / "pointworld",
        help="PointWorldStore local_dir (must already hold the flows.h5 + cameras.json)",
    )
    parser.add_argument(
        "--out", type=Path, default=None,
        help="JSON sidecar path; defaults to outputs/verify_world_frame/<uuid>__<clip>.json",
    )
    parser.add_argument("--log-level", default="INFO")
    return parser.parse_args()


def _load_camera(clip: SceneFlowClip) -> Camera:
    """Build a :class:`Camera` rescaled to ``clip.initial_rgb``'s resolution.

    PointWorld's annotation intrinsic is already valid at ``initial_rgb``'s
    resolution (both are 320x180 in every clip inspected), so ``.rescaled()``
    is a no-op here -- it is still called explicitly because it is the only
    sanctioned way to guarantee that per :meth:`Camera.rescaled`'s docstring,
    rather than relying on that equality holding by convention.
    """
    if clip.initial_rgb is None:
        raise DataError(f"clip {clip.key!r} camera {clip.camera_serial!r} has no initial_rgb")
    h, w = clip.initial_rgb.shape[:2]
    camera = Camera.from_pointworld(clip.intrinsic, clip.extrinsic, width=w, height=h)
    return camera.rescaled(w, h)


def _to_world(clip: SceneFlowClip, camera: Camera, convention: FrameConvention) -> np.ndarray:
    """Normalize ``clip.scene_flows`` (all T frames) to the world frame.

    The camera is static for the whole clip (one ``extrinsic`` per clip, not
    per frame), so the same transform applies to every frame once the
    convention is known from frame 0.
    """
    flows = clip.scene_flows.astype(np.float64)
    if convention == FrameConvention.WORLD:
        return flows
    return camera.cam_to_world(flows)


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

    Duplicated from (not imported off of) ``fpgm.geometry.convention``'s private
    ``_bilinear_sample``: that module is owned by another workstream and its
    underscore-prefixed helper isn't part of its public contract, so this script
    keeps its own small copy rather than depending on it. Callers must
    pre-filter ``uv`` to be in-bounds; this does not bounds-check.
    """
    h, w = image.shape[:2]
    u, v = uv[:, 0], uv[:, 1]
    u0 = np.clip(np.floor(u).astype(np.int64), 0, w - 1)
    v0 = np.clip(np.floor(v).astype(np.int64), 0, h - 1)
    u1 = np.clip(u0 + 1, 0, w - 1)
    v1 = np.clip(v0 + 1, 0, h - 1)
    fu = (u - u0).reshape(-1, 1)
    fv = (v - v0).reshape(-1, 1)
    img = image.astype(np.float64)
    top = img[v0, u0] * (1 - fu) + img[v0, u1] * fu
    bottom = img[v1, u0] * (1 - fu) + img[v1, u1] * fu
    return top * (1 - fv) + bottom * fv


def _cross_camera_test(
    label: str,
    world_pts_src: np.ndarray,  # (M, 3) float64, valid frame-0 world points from the source camera
    colors_src: np.ndarray,  # (M, 3) uint8, that camera's own scene_colors for the same points
    camera_dst: Camera,
    rgb_dst: np.ndarray,  # (H, W, 3) uint8
    depth_dst_mm: np.ndarray | None,  # (H, W) uint16 millimetres, or None if unavailable
) -> dict:
    """Project ``world_pts_src`` into camera B and compare colour/depth there.

    Args:
        label: Human-readable direction label, e.g. ``"A_into_B"``.
        world_pts_src: Frame-0 world points already normalized to world frame,
            restricted to points valid (visible + depth-valid) in the source
            camera's own frame 0.
        colors_src: The source camera's own recorded colour for each point --
            the "ground truth" this projection is checked against.
        camera_dst: Destination camera to project into.
        rgb_dst: Destination camera's frame-0 RGB image.
        depth_dst_mm: Destination camera's dense first-frame depth, if present.

    Returns:
        A dict report: point counts, in-bounds fraction, and median per-channel
        colour error / median depth error (``None`` where no point landed
        in-bounds, which must be surfaced as inconclusive rather than a false 0).
    """
    uv, depth = camera_dst.project(world_pts_src)
    h, w = rgb_dst.shape[:2]
    in_bounds = (uv[:, 0] >= 0) & (uv[:, 0] < w - 1) & (uv[:, 1] >= 0) & (uv[:, 1] < h - 1) & (
        depth > 0
    )
    n_points = int(world_pts_src.shape[0])
    frac_in_bounds = float(np.mean(in_bounds)) if n_points else 0.0
    report: dict = {
        "label": label,
        "n_points": n_points,
        "n_in_bounds": int(np.count_nonzero(in_bounds)),
        "frac_in_bounds": frac_in_bounds,
        "median_abs_color_diff_per_channel": None,
        "median_abs_color_diff_mean": None,
        "median_abs_depth_error_mm": None,
        "n_depth_compared": 0,
    }
    if not np.any(in_bounds):
        return report

    sampled = _bilinear_sample(rgb_dst, uv[in_bounds])
    reference = colors_src[in_bounds].astype(np.float64)
    diff = np.abs(sampled - reference)
    median_per_channel = np.median(diff, axis=0).tolist()
    report["median_abs_color_diff_per_channel"] = median_per_channel
    report["median_abs_color_diff_mean"] = float(np.mean(median_per_channel))

    if depth_dst_mm is not None:
        u_px = np.clip(np.round(uv[in_bounds, 0]).astype(np.int64), 0, w - 1)
        v_px = np.clip(np.round(uv[in_bounds, 1]).astype(np.int64), 0, h - 1)
        depth_dst_at_px_mm = depth_dst_mm[v_px, u_px].astype(np.float64)
        depth_valid = depth_dst_at_px_mm > 0  # 0 = no dense-depth sample at that pixel
        if np.any(depth_valid):
            reproj_depth_mm = depth[in_bounds][depth_valid] * 1000.0
            abs_err = np.abs(reproj_depth_mm - depth_dst_at_px_mm[depth_valid])
            report["median_abs_depth_error_mm"] = float(np.median(abs_err))
            report["n_depth_compared"] = int(np.count_nonzero(depth_valid))

    return report


def _rigid_transform_check(
    clip_a: SceneFlowClip, clip_b: SceneFlowClip, cameras_json: dict
) -> dict:
    """Compare the A->B rigid transform implied by flows.h5 vs by cameras.json.

    Both ``clip.extrinsic`` (from ``*_flows.h5``) and ``optimized_extrinsics``
    (from ``<uuid>_cameras.json``) claim to be the same world->camera SE3 per
    serial; this checks they actually agree by composing each source's own A/B
    pair into a relative transform and comparing those, rather than comparing
    the raw matrices directly (which would also catch a difference in *which*
    world frame each source roots its extrinsic in -- exactly the failure mode
    this whole script exists to catch).
    """
    rel_flows = clip_b.extrinsic @ invert_se3(clip_a.extrinsic)

    opt_a = np.asarray(
        cameras_json[clip_a.camera_serial]["optimized_extrinsics"], dtype=np.float64
    )
    opt_b = np.asarray(
        cameras_json[clip_b.camera_serial]["optimized_extrinsics"], dtype=np.float64
    )
    rel_json = opt_b @ invert_se3(opt_a)

    translation_error_m = float(np.linalg.norm(rel_flows[:3, 3] - rel_json[:3, 3]))
    r_err = rel_flows[:3, :3].T @ rel_json[:3, :3]
    rotation_error_deg = float(np.degrees(np.linalg.norm(matrix_to_rotvec(r_err))))
    return {
        "translation_error_m": translation_error_m,
        "rotation_error_deg": rotation_error_deg,
    }


def _print_verdict(name: str, passed: bool, detail: str) -> None:
    tag = "PASS" if passed else "FAIL"
    print(f"[{tag}] {name}: {detail}")


def main() -> int:
    args = parse_args()
    setup_logging(args.log_level)

    store = PointWorldStore(local_dir=args.data_dir)
    flows_path = store.flows_path(args.episode)
    with FlowsReader(flows_path, episode_uuid=args.episode) as reader:
        clip_a = reader.read_clip(args.clip, args.camera_a)
        clip_b = reader.read_clip(args.clip, args.camera_b)

    camera_a = _load_camera(clip_a)
    camera_b = _load_camera(clip_b)

    report: dict = {
        "episode": args.episode,
        "clip": args.clip,
        "camera_a": args.camera_a,
        "camera_b": args.camera_b,
        "n_points_a": clip_a.n_points,
        "n_points_b": clip_b.n_points,
        "n_frames": clip_a.n_frames,
    }

    # -- 1. convention detection ------------------------------------------------
    cfg = ConventionConfig()
    conventions: dict[str, FrameConvention | None] = {}
    convention_report: dict = {}
    for label, clip, camera in (("a", clip_a, camera_a), ("b", clip_b, camera_b)):
        try:
            det = detect_scene_flow_convention(
                camera, clip.scene_flows[0], clip.scene_colors[0], clip.initial_rgb, cfg
            )
            conventions[label] = det.convention
            convention_report[label] = {
                "convention": det.convention.value,
                "score_world": det.score_world,
                "score_camera": det.score_camera,
                "margin": det.margin,
                "frac_in_bounds": det.frac_in_bounds,
            }
        except (AmbiguousConventionError, CameraCalibrationMismatchError) as exc:
            conventions[label] = None
            convention_report[label] = {"error": str(exc)}
    report["convention"] = convention_report

    convention_ok = conventions["a"] is not None and conventions["b"] is not None
    if convention_ok:
        detail = f"a={conventions['a'].value}, b={conventions['b'].value}"
    else:
        detail = f"detection failed: {convention_report}"
    _print_verdict("convention_detection", convention_ok, detail)

    if not convention_ok:
        report["verdict"] = "INCONCLUSIVE"
        report["verdict_detail"] = "convention detection failed for at least one camera"
        _write_report(args, report)
        print("\nOVERALL: INCONCLUSIVE -- cannot test cross-camera agreement without a "
              "decided convention for both cameras.")
        return 1

    # -- 2 & 3. cross-camera colour + depth agreement ----------------------------
    world_a = _to_world(clip_a, camera_a, conventions["a"])
    world_b = _to_world(clip_b, camera_b, conventions["b"])

    valid_a0 = clip_a.scene_visibility[0] & clip_a.scene_depth_valid[0]
    valid_b0 = clip_b.scene_visibility[0] & clip_b.scene_depth_valid[0]

    a_into_b = _cross_camera_test(
        "A_into_B", world_a[0][valid_a0], clip_a.scene_colors[0][valid_a0],
        camera_b, clip_b.initial_rgb, clip_b.initial_depth,
    )
    b_into_a = _cross_camera_test(
        "B_into_A", world_b[0][valid_b0], clip_b.scene_colors[0][valid_b0],
        camera_a, clip_a.initial_rgb, clip_a.initial_depth,
    )
    report["cross_camera"] = {"a_into_b": a_into_b, "b_into_a": b_into_a}

    color_checks = []
    depth_checks = []
    for direction in (a_into_b, b_into_a):
        color_ok = (
            direction["frac_in_bounds"] >= _PASS_MIN_FRAC_IN_BOUNDS
            and direction["median_abs_color_diff_mean"] is not None
            and direction["median_abs_color_diff_mean"] <= _PASS_COLOR_DIFF_PER_CHANNEL
        )
        color_checks.append(color_ok)
        _print_verdict(
            f"color_agreement[{direction['label']}]",
            color_ok,
            f"frac_in_bounds={direction['frac_in_bounds']:.3f}, "
            f"median_abs_diff_per_channel={direction['median_abs_color_diff_per_channel']}",
        )
        depth_ok = (
            direction["n_depth_compared"] > 0
            and direction["median_abs_depth_error_mm"] is not None
            and direction["median_abs_depth_error_mm"] <= _PASS_DEPTH_ERROR_MM
        )
        depth_checks.append(depth_ok)
        _print_verdict(
            f"depth_agreement[{direction['label']}]",
            depth_ok,
            f"n_compared={direction['n_depth_compared']}, "
            f"median_abs_depth_error_mm={direction['median_abs_depth_error_mm']}",
        )

    # -- 4. rigid transform vs cameras.json --------------------------------------
    cameras_json = store.load_camera_calibration(args.episode)
    rigid = _rigid_transform_check(clip_a, clip_b, cameras_json)
    report["rigid_transform_vs_cameras_json"] = rigid
    rigid_ok = (
        rigid["translation_error_m"] <= _PASS_TRANSLATION_ERROR_M
        and rigid["rotation_error_deg"] <= _PASS_ROTATION_ERROR_DEG
    )
    _print_verdict(
        "rigid_transform_vs_cameras_json",
        rigid_ok,
        f"translation_error_m={rigid['translation_error_m']:.6f}, "
        f"rotation_error_deg={rigid['rotation_error_deg']:.6f}",
    )

    overall_ok = convention_ok and all(color_checks) and all(depth_checks) and rigid_ok
    report["verdict"] = "PASS" if overall_ok else "FAIL"
    _write_report(args, report)

    print(f"\nOVERALL: {'PASS' if overall_ok else 'FAIL'} -- cameras {args.camera_a} and "
          f"{args.camera_b} of clip {args.clip} "
          f"{'agree on' if overall_ok else 'DO NOT clearly agree on'} a common world frame.")
    if not overall_ok:
        print("See per-test numbers above and the JSON sidecar before building anything "
              "on a single canonical world frame assumption.")
    return 0 if overall_ok else 1


def _write_report(args: argparse.Namespace, report: dict) -> None:
    out = args.out
    if out is None:
        out_dir = REPO_ROOT / "outputs" / "verify_world_frame"
        out_dir.mkdir(parents=True, exist_ok=True)
        out = out_dir / f"{args.episode}__{args.clip.replace(':', '-')}.json"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(report, indent=2))
    logger.info("wrote %s", out)


if __name__ == "__main__":
    raise SystemExit(main())