scripts/run_object_pipeline.py

#!/usr/bin/env python
"""Staged object reconstruction + interaction pipeline (stages 1-5), the debugging backbone.

    PYTHONPATH=src python scripts/run_object_pipeline.py \\
        --episode <uuid> --camera ext1 --clip 10:21 --stages mask,mesh,align,act,render

Five stages, each independently re-runnable from the previous stage's artefacts
on disk (see :mod:`fpgm.objects.types`'s module docstring):

    1. mask   : SAM 3.1 text-prompted segmentation -> tight RGBA crop + the
                frame-0 mask (frame 0 is the only frame this dataset gives us
                dense depth for, via ``initial_depth``/``initial_rgb``).
    2. mesh   : crop -> mesh, either TRELLIS or SAM 3D Objects (both externally
                generated, passed via ``--glb``) or a proxy fitted to the
                observed point cloud.
    3. align  : mask x depth -> world point cloud -> scale + 6-DOF mesh pose.
    4. act    : robot FK -> per-frame object pose (free / pushed / grasped).
    5. render : robot + object + point cloud, composited in the world frame,
                from both the recorded camera and a synthetic orbiting one.

Each stage writes into ``outputs/<uuid>/objects/<n>_<stage>/`` and a later
stage always *loads* its inputs from there rather than recomputing them --
run e.g. only ``--stages render`` once ``mask``, ``mesh``, ``align`` and
``act`` already have artefacts on disk. A missing upstream artefact fails with
the exact command to produce it, not a stack trace.

Example, once ``mask``/``mesh``/``align``/``act`` have already been run:

    PYTHONPATH=src python scripts/run_object_pipeline.py \\
        --episode <uuid> --camera ext1 --clip 10:21 --stages render --push-gain 8.0
"""

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass
from pathlib import Path

import cv2
import h5py
import numpy as np
import trimesh
from PIL import Image

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

from fpgm.data.droid_raw import cam2world_vector_to_world2cam, read_mp4_properties  # noqa: E402
from fpgm.data.pointworld import FlowsReader, verify_annotation_stride  # noqa: E402
from fpgm.geometry.camera import Camera  # noqa: E402
from fpgm.objects.types import (  # noqa: E402
    InteractionState,
    MeshSource,
    ObjectAlignment,
    ObjectMesh,
    ObjectTrajectory,
    StageArtifacts,
)
from fpgm.pipeline.frames import ClipFrameSource  # noqa: E402
from fpgm.types import ClipTiming, DataError, SceneFlowClip  # noqa: E402
from fpgm.utils.io import ensure_dir  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402
from fpgm.viz.overlays import VideoWriter  # noqa: E402

logger = get_logger("run_object_pipeline")

STAGES: tuple[str, ...] = ("mask", "mesh", "align", "act", "render")
STAGE_INDEX: dict[str, int] = {s: i + 1 for i, s in enumerate(STAGES)}

_DEFAULT_URDF = (
    REPO_ROOT
    / "third_party"
    / "robot_description"
    / "pointworld_franka_robotiq_2f85"
    / "franka_panda_robotiq_2f85.urdf"
)

#: DROID's nominal 15 Hz control rate. The public raw release does not embed
#: the real one in trajectory.h5 (see fpgm.data.episode's own note on this),
#: so this fallback is the *expected* path here, not a last resort.
_TRAJECTORY_FPS = 15.0

#: Neutral background for the free-camera video, which has no real video plate
#: behind it (the recorded camera never saw this synthetic viewpoint).
_FREECAM_BG = (40, 40, 40)


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("--episode", required=True, help="DROID episode uuid")
    p.add_argument("--camera", default="ext1", choices=["ext1", "ext2"])
    p.add_argument(
        "--clip", required=True, help='PointWorld scene-flow clip key, e.g. "10:21"'
    )
    p.add_argument(
        "--stages",
        default=",".join(STAGES),
        help=f"comma-separated stages to run, in order; any subset of {STAGES}",
    )
    p.add_argument(
        "--mesh-source",
        default="proxy_box",
        choices=["trellis", "sam3d", "proxy_box", "proxy_hull"],
        help="mesh stage source; defaults to a fitted proxy box until a real "
        "generator's output is supplied. 'sam3d' is facebook/sam-3d-objects "
        "(see scripts/sam3d_generate.py) -- as of this writing its checkpoint "
        "access is pending review, so this path is untested end-to-end; the "
        "code is complete and should work once a --glb is produced.",
    )
    p.add_argument(
        "--glb", type=Path, default=None,
        help="externally generated mesh (from TRELLIS or SAM 3D Objects, produced by "
        "scripts/trellis_generate.py or scripts/sam3d_generate.py respectively); "
        "required with --mesh-source trellis or --mesh-source sam3d",
    )
    p.add_argument(
        "--push-gain", type=float, default=4.0,
        help="visualisation-only amplification of pushed motion (act stage); 1.0 = physical",
    )
    p.add_argument(
        "--outputs", type=Path, default=None,
        help="output directory (default: outputs/<episode-uuid>)",
    )
    p.add_argument("--urdf", type=Path, default=_DEFAULT_URDF, help="robot URDF path")
    p.add_argument(
        "--device", default="cuda", help="torch device for SAM 3.1 (mask stage only)"
    )
    p.add_argument("--seed", type=int, default=0, help="RNG seed (mesh-fit ICP, sampling)")
    p.add_argument(
        "--freecam-distance", type=float, default=None,
        help="orbit radius, metres (render stage); default: 3x the object's extent",
    )
    p.add_argument(
        "--freecam-elevation", type=float, default=20.0,
        help="orbit elevation above horizontal, degrees (render stage)",
    )
    return p.parse_args()


# --------------------------------------------------------------------------- #
# Shared data access -- every stage re-derives what it needs from raw episode
# data rather than trusting an earlier stage's in-memory state, so each one
# genuinely stands alone (see the module docstring).
# --------------------------------------------------------------------------- #


def _episode_dir(uuid: str) -> Path:
    return REPO_ROOT / "data" / "droid_raw" / uuid


def _flows_path(uuid: str) -> Path:
    return REPO_ROOT / "data" / "pointworld" / "droid" / "flows-fs-optimized" / f"{uuid}_flows.h5"


def _load_metadata(uuid: str) -> dict:
    path = _episode_dir(uuid) / "metadata.json"
    if not path.exists():
        raise DataError(
            f"no local DROID metadata for episode {uuid!r} at {path}; fetch it first with "
            f"`python scripts/download_droid_episodes.py --episode {uuid}`."
        )
    return json.loads(path.read_text())


def _camera_serial(metadata: dict, role: str) -> str:
    serial = metadata.get(f"{role}_cam_serial")
    if serial is None:
        raise DataError(f"episode metadata has no {role}_cam_serial entry")
    return serial


@dataclass
class ClipContext:
    """Everything derived from raw episode data that every stage needs."""

    clip: SceneFlowClip
    camera_annot: Camera  # native annotation resolution -- matches initial_depth/initial_rgb
    mp4_path: Path
    serial: str
    #: Clock reconciliation for this clip, including the *measured* annotation
    #: clip-frame -> trajectory/video-row stride. See :class:`fpgm.types.ClipTiming`.
    timing: ClipTiming


def _build_timing(uuid: str, clip: SceneFlowClip) -> ClipTiming:
    """Measure ``clip``'s annotation-frame -> trajectory-row stride from real data.

    Hardcoding the previously-discovered stride of 2 would repeat exactly the
    mistake that produced the original bug -- see
    :func:`fpgm.data.pointworld.verify_annotation_stride` and
    :class:`fpgm.types.ClipTiming`'s docstring -- so every clip proves it against
    its own ``joint_positions`` and this episode's own ``trajectory.h5``.
    """
    if clip.joint_positions is None:
        raise DataError(
            f"clip {clip.key} has no joint_positions; cannot verify its "
            "annotation_stride against trajectory.h5"
        )
    traj_path = _load_trajectory_h5(uuid)
    with h5py.File(traj_path, "r") as f:
        trajectory_joint_positions = np.asarray(f["observation/robot_state/joint_positions"])
    stride = verify_annotation_stride(clip.joint_positions, trajectory_joint_positions, clip.start)
    return ClipTiming(
        clip_start_frame=clip.start,
        clip_end_frame=clip.end,
        trajectory_fps=_TRAJECTORY_FPS,
        annotation_stride=stride,
    )


def _strided_rows(start: int, n_frames: int, stride: int, available: int, what: str) -> np.ndarray:
    """Row/frame indices ``start*stride + stride*[0..n_frames)``, clipped to ``available``.

    A mapped range is now up to ``stride`` times wider than the clip's own frame
    count, so running off the end of a source array (``trajectory.h5``, the mp4) is
    a real possibility, not just defensive boilerplate -- e.g. clip ``"50:61"`` maps
    to trajectory/video rows ``[100, 122)``, which is within a 127-frame episode
    but only just. Rather than raising, the axis is truncated to what actually
    exists and a warning is logged, matching how :class:`~fpgm.pipeline.frames.
    ClipFrameSource` already handles a video shorter than the requested range.
    """
    rows = start * stride + stride * np.arange(n_frames)
    in_range = rows < available
    if not np.all(in_range):
        logger.warning(
            "%s: mapped range needs rows up to %d but only %d available "
            "(stride=%d); truncating from %d to %d frames",
            what,
            int(rows[-1]),
            available,
            stride,
            n_frames,
            int(in_range.sum()),
        )
        rows = rows[in_range]
    return rows


def _load_clip_context(uuid: str, camera_role: str, clip_key: str) -> ClipContext:
    metadata = _load_metadata(uuid)
    serial = _camera_serial(metadata, camera_role)
    mp4_path = _episode_dir(uuid) / "recordings" / "MP4" / f"{serial}.mp4"
    if not mp4_path.exists():
        raise DataError(f"mp4 not found for camera {serial!r}: {mp4_path}")

    flows_path = _flows_path(uuid)
    if not flows_path.exists():
        raise DataError(f"PointWorld flows.h5 not found for episode {uuid!r} at {flows_path}")
    with FlowsReader(flows_path, episode_uuid=uuid) as reader:
        available = reader.clip_keys()
        if clip_key not in available:
            raise DataError(
                f"clip {clip_key!r} not found in {flows_path.name}; available clips: {available}"
            )
        clip = reader.read_clip(clip_key, serial)

    if clip.initial_rgb is None:
        raise DataError(f"clip {clip_key} has no initial_rgb to calibrate the camera against")
    annot_h, annot_w = clip.initial_rgb.shape[:2]
    camera_annot = Camera.from_pointworld(
        clip.intrinsic, clip.extrinsic, width=annot_w, height=annot_h
    )
    timing = _build_timing(uuid, clip)
    return ClipContext(
        clip=clip, camera_annot=camera_annot, mp4_path=mp4_path, serial=serial, timing=timing
    )


def _load_initial_depth(uuid: str, clip_key: str, serial: str) -> np.ndarray:
    """``(H, W)`` uint16 mm depth, present per-camera in flows.h5 but not (yet)
    exposed on :class:`~fpgm.types.SceneFlowClip`, so it is read directly here."""
    with h5py.File(_flows_path(uuid), "r") as f:
        return np.asarray(f[clip_key][f"camera_{serial}_ext"]["initial_depth"])


class _DepthRgbAdapter:
    """Satisfies :class:`fpgm.objects.align.DepthRgbSource` without widening
    the shared :class:`~fpgm.types.SceneFlowClip` contract for one stage."""

    def __init__(self, initial_depth: np.ndarray, initial_rgb: np.ndarray | None) -> None:
        self.initial_depth = initial_depth
        self.initial_rgb = initial_rgb


def _load_trajectory_h5(uuid: str) -> Path:
    path = _episode_dir(uuid) / "trajectory.h5"
    if not path.exists():
        raise DataError(f"trajectory.h5 not found: {path}")
    return path


# --------------------------------------------------------------------------- #
# Stage plumbing: directories, and the "produce it first" failure message.
# --------------------------------------------------------------------------- #


def stage_dir(outputs_root: Path, stage: str) -> Path:
    return outputs_root / "objects" / f"{STAGE_INDEX[stage]}_{stage}"


def _cmd_for(args: argparse.Namespace, stages: str) -> str:
    parts = [
        "PYTHONPATH=src", "python", "scripts/run_object_pipeline.py",
        "--episode", args.episode, "--camera", args.camera, "--clip", args.clip,
        "--stages", stages,
    ]
    if args.mesh_source != "proxy_box":
        parts += ["--mesh-source", args.mesh_source]
    if args.glb:
        parts += ["--glb", str(args.glb)]
    return " ".join(parts)


def _require(path: Path, stage: str, args: argparse.Namespace, upstream_stages: str) -> Path:
    if not path.exists():
        raise DataError(
            f"{stage} stage: missing required input {path}.\n"
            f"Produce it first with:\n    {_cmd_for(args, upstream_stages)}"
        )
    return path


# --------------------------------------------------------------------------- #
# Stage 1: mask
# --------------------------------------------------------------------------- #


def run_mask(args: argparse.Namespace, outputs_root: Path) -> StageArtifacts:
    from fpgm.config import SegmentationConfig
    from fpgm.objects import crop
    from fpgm.pipeline.velocity_pipeline import _prompt_variants
    from fpgm.prompting.task_prompt import TaskPromptDeriver
    from fpgm.segmentation.sam3 import Sam3VideoSegmenter

    out_dir = ensure_dir(stage_dir(outputs_root, "mask"))
    ctx = _load_clip_context(args.episode, args.camera, args.clip)
    metadata = _load_metadata(args.episode)
    task = metadata.get("current_task", "")
    # Reuse the velocity pipeline's candidate expansion rather than the bare head
    # noun: on this data SAM 3.1 returns nothing for "brick" but scores 0.93 for
    # "blue block", so the colour-qualified variants are the normal path.
    derived = TaskPromptDeriver().derive(task)
    candidates: list[str] = []
    for phrase in [derived.primary, *derived.alternatives]:
        for variant in _prompt_variants(phrase):
            if variant and variant not in candidates:
                candidates.append(variant)
    logger.info("clip %s: task=%r -> prompt candidates %s", args.clip, task, candidates)

    seg_cfg = SegmentationConfig()
    segmenter = Sam3VideoSegmenter(seg_cfg, device=args.device)

    # The clip's declared [start, end) is in the annotation clock, not the
    # video's -- map through ctx.timing (measured annotation_stride) so SAM sees
    # exactly the video frames the annotations describe, stride included.
    video_start = int(ctx.timing.clip_frame_to_video_frame(0))
    video_end = int(ctx.timing.clip_frame_to_video_frame(ctx.timing.n_frames))
    with ClipFrameSource(
        str(ctx.mp4_path), video_start, video_end, stride=ctx.timing.annotation_stride
    ) as frames:
        width, height = frames.resolution
        frame0_bgr = frames.read(0)
        cv2.imwrite(str(out_dir / "frame0_bgr.png"), frame0_bgr)

        masklets = None
        obj_id: int | None = None
        used_prompt = ""
        for prompt in candidates:
            segmenter.start_session(str(frames.frame_dir))
            try:
                segmenter.add_text_prompt(seg_cfg.prompt_frame_idx, prompt)
                found = segmenter.collect_masklets()
                if not found:
                    logger.info("prompt %r matched nothing; trying next candidate", prompt)
                    continue
                obj_id = segmenter.select_object(found)
                masklets, used_prompt = found, prompt
                break
            finally:
                segmenter.close_session()

        if masklets is None or obj_id is None:
            raise DataError(f"mask stage: SAM 3.1 found nothing for any of {candidates}")

        masklet = masklets[obj_id]
        if 0 not in masklet.frames or not masklet.frames[0].any():
            raise DataError(
                f"mask stage: object mask is empty on frame 0 of clip {args.clip}; frame 0 "
                "is the only frame with dense depth (initial_depth), so reconstruction "
                "cannot proceed without it -- try a different --clip"
            )
        mask_frame0 = masklet.frames[0]

        best_frame = crop.select_best_frame(masklet, prefer="largest")
        best_frame_bgr = frames.read(best_frame)
        obj_crop = crop.crop_from_mask(
            best_frame_bgr, masklet.frames[best_frame], frame_idx=best_frame, prompt=used_prompt
        )
        artifacts = crop.save_debug(obj_crop, out_dir, frame_bgr=best_frame_bgr)

    np.savez_compressed(out_dir / "mask_frame0.npz", mask=mask_frame0)
    meta = {
        "episode": args.episode,
        "camera": args.camera,
        "camera_serial": ctx.serial,
        "clip": args.clip,
        "prompt": used_prompt,
        "obj_id": int(obj_id),
        "video_resolution": [width, height],
        "best_frame": int(best_frame),
        "n_masklets": len(masklets),
    }
    (out_dir / "meta.json").write_text(json.dumps(meta, indent=2))

    artifacts.add("mask_frame0", out_dir / "mask_frame0.npz")
    artifacts.add("frame0_bgr", out_dir / "frame0_bgr.png")
    artifacts.add("meta", out_dir / "meta.json")
    artifacts.stats.update(meta)
    return artifacts


def _load_mask_meta(outputs_root: Path) -> dict:
    return json.loads((stage_dir(outputs_root, "mask") / "meta.json").read_text())


def _load_mask_frame0(outputs_root: Path) -> np.ndarray:
    with np.load(stage_dir(outputs_root, "mask") / "mask_frame0.npz") as npz:
        return npz["mask"].astype(bool)


# --------------------------------------------------------------------------- #
# Stage 2: mesh
# --------------------------------------------------------------------------- #


def _visual_to_mesh_fields(visual: trimesh.visual.base.Visuals) -> tuple[
    np.ndarray | None, np.ndarray | None, np.ndarray | None
]:
    """Pull ``(vertex_colors, texture, uv)`` out of a loaded trimesh ``visual``.

    A textured GLB (TRELLIS/SAM 3D output) loads as ``TextureVisuals``, which
    has no ``vertex_colors`` attribute at all -- ``getattr(visual,
    "vertex_colors", None)`` silently returns ``None`` for it, which used to
    make this loader throw away both the texture AND any colour on every
    textured mesh (SAM3D/TRELLIS bricks rendered flat grey as a result). The
    ``kind == "texture"`` branch below is what actually keeps the texture;
    ``to_color().vertex_colors`` in the fallback branch is only there so a
    non-textured mesh (proxy hull/box, or a bare vertex-coloured GLB) still
    keeps *some* colour instead of silently losing it the same way.
    """
    vertex_colors: np.ndarray | None = None
    texture: np.ndarray | None = None
    uv: np.ndarray | None = None
    if getattr(visual, "kind", None) == "texture":
        uv = np.asarray(visual.uv, dtype=np.float32) if visual.uv is not None else None
        base_color = getattr(visual.material, "baseColorTexture", None)
        if base_color is not None:
            texture = np.asarray(base_color.convert("RGB"), dtype=np.uint8)
        if uv is None or texture is None:
            # A "texture" visual with no actual image/uv (e.g. material-only,
            # no baked pixels) -- fall back to flat colour rather than pass
            # along a half-populated (uv, texture) pair that render code would
            # have to special-case.
            uv, texture = None, None
            colors = visual.to_color().vertex_colors
            if colors is not None and len(colors):
                vertex_colors = np.asarray(colors)[:, :3].astype(np.uint8)
    else:
        colors = visual.to_color().vertex_colors
        if colors is not None and len(colors):
            vertex_colors = np.asarray(colors)[:, :3].astype(np.uint8)
    return vertex_colors, texture, uv


def _save_mesh(mesh: ObjectMesh, out_dir: Path) -> None:
    # process=False everywhere in this function: trimesh's default vertex
    # merging/welding would change which row of `mesh.uv` corresponds to which
    # vertex, silently misaligning the texture from the geometry it was baked
    # onto.
    if mesh.uv is not None and mesh.texture is not None:
        visual = trimesh.visual.TextureVisuals(
            uv=np.asarray(mesh.uv, dtype=np.float32),
            image=Image.fromarray(np.asarray(mesh.texture, dtype=np.uint8)),
        )
        tri = trimesh.Trimesh(
            vertices=np.asarray(mesh.vertices, dtype=np.float64),
            faces=np.asarray(mesh.faces, dtype=np.int64),
            visual=visual,
            process=False,
        )
    else:
        tri = trimesh.Trimesh(
            vertices=np.asarray(mesh.vertices, dtype=np.float64),
            faces=np.asarray(mesh.faces, dtype=np.int64),
            process=False,
        )
        if mesh.vertex_colors is not None:
            tri.visual.vertex_colors = mesh.vertex_colors
    tri.export(str(out_dir / "mesh.glb"))

    meta: dict = {"source": mesh.source.value}
    npz_payload: dict[str, np.ndarray] = {}
    for key, value in mesh.metadata.items():
        if isinstance(value, np.ndarray):
            npz_payload[key] = value
        else:
            meta[key] = value
    if npz_payload:
        np.savez_compressed(out_dir / "mesh_metadata.npz", **npz_payload)
    (out_dir / "mesh_meta.json").write_text(json.dumps(meta, indent=2))


def _load_mesh(mesh_dir: Path) -> ObjectMesh:
    meta = json.loads((mesh_dir / "mesh_meta.json").read_text())
    # force="mesh" is not the problem here (it preserves TextureVisuals, see
    # trimesh/visual/objects.py); process=False is what matters, for the same
    # vertex/UV-correspondence reason as in _save_mesh.
    loaded = trimesh.load(str(mesh_dir / "mesh.glb"), force="mesh", process=False)
    vertex_colors, texture, uv = _visual_to_mesh_fields(loaded.visual)

    metadata = {k: v for k, v in meta.items() if k != "source"}
    npz_path = mesh_dir / "mesh_metadata.npz"
    if npz_path.exists():
        with np.load(npz_path) as npz:
            metadata.update({k: npz[k] for k in npz.files})

    return ObjectMesh(
        vertices=np.asarray(loaded.vertices, dtype=np.float64),
        faces=np.asarray(loaded.faces, dtype=np.int64),
        source=MeshSource(meta["source"]),
        vertex_colors=vertex_colors,
        texture=texture,
        uv=uv,
        metadata=metadata,
    )


def run_mesh(args: argparse.Namespace, outputs_root: Path) -> StageArtifacts:
    from fpgm.objects.align import object_point_cloud
    from fpgm.objects.proxy import convex_hull_mesh, oriented_box_mesh

    out_dir = ensure_dir(stage_dir(outputs_root, "mesh"))
    mask_dir = stage_dir(outputs_root, "mask")
    _require(mask_dir / "meta.json", "mesh", args, "mask")
    _require(mask_dir / "mask_frame0.npz", "mesh", args, "mask")
    meta = _load_mask_meta(outputs_root)

    if args.mesh_source in ("trellis", "sam3d"):
        # Both externally-generated sources share the exact same contract here:
        # a standalone script (scripts/trellis_generate.py or
        # scripts/sam3d_generate.py), run beforehand in that generator's own
        # conda env (trellis2 / sam3d-objects respectively -- see each
        # script's module docstring), writes a textured/vertex-colored .glb;
        # this stage only loads it. Generation is never shelled out to from
        # here, matching how --mesh-source trellis has always worked.
        #
        # NOTE(sam3d): scripts/sam3d_generate.py was validated end-to-end
        # against the sam-3d-objects repo's own sample image (2026-07-30; see
        # logs/sam3d_setup/09_smoketest_generate.log and
        # 11_smoketest_textured.log) and this loading branch was confirmed to
        # accept its .glb output (identical trimesh contract to the trellis
        # path). Not yet exercised with a real DROID episode crop through
        # this actual CLI flag end-to-end -- only the sam3d_generate.py step
        # itself and this loader were checked independently.
        if args.glb is None:
            raise DataError(
                f"mesh stage: --mesh-source {args.mesh_source} requires --glb "
                f"<path to {'TRELLIS' if args.mesh_source == 'trellis' else 'SAM 3D Objects'} output>"
            )
        if not args.glb.exists():
            raise DataError(f"mesh stage: --glb path does not exist: {args.glb}")
        # process=False: see _save_mesh's docstring -- welding vertices here
        # would desync `uv` (per-vertex) from the geometry it was baked onto.
        loaded = trimesh.load(str(args.glb), force="mesh", process=False)
        vertex_colors, texture, uv = _visual_to_mesh_fields(loaded.visual)
        mesh = ObjectMesh(
            vertices=np.asarray(loaded.vertices, dtype=np.float64),
            faces=np.asarray(loaded.faces, dtype=np.int64),
            source=MeshSource.TRELLIS if args.mesh_source == "trellis" else MeshSource.SAM3D,
            vertex_colors=vertex_colors,
            texture=texture,
            uv=uv,
            metadata={"source_glb": str(args.glb)},
        )
    else:
        mask_full = _load_mask_frame0(outputs_root)
        video_w, video_h = meta["video_resolution"]
        if mask_full.shape != (video_h, video_w):
            raise DataError(
                f"mesh stage: mask shape {mask_full.shape} != recorded video resolution "
                f"{(video_h, video_w)}"
            )
        ctx = _load_clip_context(args.episode, args.camera, args.clip)
        depth = _load_initial_depth(args.episode, args.clip, meta["camera_serial"])
        adapter = _DepthRgbAdapter(depth, ctx.clip.initial_rgb)
        points_world, _colors = object_point_cloud(
            adapter, ctx.camera_annot, mask_full, (video_h, video_w)
        )
        builder = oriented_box_mesh if args.mesh_source == "proxy_box" else convex_hull_mesh
        mesh = builder(points_world)

    _save_mesh(mesh, out_dir)
    stats = {
        "source": mesh.source.value,
        "n_vertices": int(mesh.vertices.shape[0]),
        "n_faces": int(mesh.faces.shape[0]),
        "extent_m": mesh.extent.tolist(),
    }
    (out_dir / "mesh_stats.json").write_text(json.dumps(stats, indent=2))

    artifacts = StageArtifacts(stage="mesh", directory=out_dir, stats=stats)
    artifacts.add("mesh_glb", out_dir / "mesh.glb")
    artifacts.add("mesh_meta", out_dir / "mesh_meta.json")
    artifacts.add("mesh_stats", out_dir / "mesh_stats.json")
    return artifacts


# --------------------------------------------------------------------------- #
# Stage 3: align
# --------------------------------------------------------------------------- #


def _load_alignment(align_dir: Path) -> ObjectAlignment:
    with np.load(align_dir / "alignment.npz") as npz:
        transform, points_world, colors = npz["transform"], npz["points_world"], npz["colors"]
    stats = json.loads((align_dir / "align_stats.json").read_text())
    return ObjectAlignment(
        transform=transform,
        scale=stats["scale"],
        points_world=points_world,
        point_colors=colors,
        rmse_m=stats["rmse_m"],
        inlier_fraction=stats["inlier_fraction"],
        frame_idx=0,
    )


def run_align(args: argparse.Namespace, outputs_root: Path) -> StageArtifacts:
    from fpgm.objects import align as align_module
    from fpgm.objects.align import fit_mesh_to_points, object_point_cloud

    out_dir = ensure_dir(stage_dir(outputs_root, "align"))
    mask_dir = stage_dir(outputs_root, "mask")
    mesh_dir = stage_dir(outputs_root, "mesh")
    _require(mask_dir / "meta.json", "align", args, "mask")
    _require(mesh_dir / "mesh.glb", "align", args, "mask,mesh")

    meta = _load_mask_meta(outputs_root)
    mask_full = _load_mask_frame0(outputs_root)
    mesh = _load_mesh(mesh_dir)

    ctx = _load_clip_context(args.episode, args.camera, args.clip)
    depth = _load_initial_depth(args.episode, args.clip, meta["camera_serial"])
    adapter = _DepthRgbAdapter(depth, ctx.clip.initial_rgb)
    video_w, video_h = meta["video_resolution"]
    points_world, colors = object_point_cloud(
        adapter, ctx.camera_annot, mask_full, (video_h, video_w)
    )

    alignment = fit_mesh_to_points(
        mesh, points_world, allow_scale=True, refine_icp=True,
        rng=np.random.default_rng(args.seed),
    )
    alignment.frame_idx = 0

    np.savez_compressed(
        out_dir / "alignment.npz",
        transform=alignment.transform, points_world=alignment.points_world, colors=colors,
    )
    stats = {
        "scale": alignment.scale,
        "rmse_m": alignment.rmse_m,
        "inlier_fraction": alignment.inlier_fraction,
        "n_points": int(points_world.shape[0]),
        "position_m": alignment.position.tolist(),
    }
    (out_dir / "align_stats.json").write_text(json.dumps(stats, indent=2))

    frame0_path = mask_dir / "frame0_bgr.png"
    align_artifacts_files = {}
    if frame0_path.exists():
        frame0_bgr = cv2.imread(str(frame0_path))
        camera_video = ctx.camera_annot.rescaled(video_w, video_h)
        debug = align_module.save_debug(alignment, camera_video, frame0_bgr, out_dir, mesh=mesh)
        align_artifacts_files.update(debug.files)

    artifacts = StageArtifacts(
        stage="align", directory=out_dir, stats=stats, files=align_artifacts_files
    )
    artifacts.add("alignment", out_dir / "alignment.npz")
    artifacts.add("align_stats", out_dir / "align_stats.json")
    return artifacts


# --------------------------------------------------------------------------- #
# Stage 4: act
# --------------------------------------------------------------------------- #


def _load_trajectory(act_dir: Path) -> ObjectTrajectory:
    with np.load(act_dir / "trajectory.npz") as npz:
        timestamps, transforms = npz["timestamps"], npz["transforms"]
        gripper_distance = npz["gripper_distance_m"]
    states = [InteractionState(s) for s in json.loads((act_dir / "states.json").read_text())]
    stats = json.loads((act_dir / "act_stats.json").read_text())
    return ObjectTrajectory(
        timestamps=timestamps, transforms=transforms, states=states,
        gripper_distance_m=gripper_distance, push_gain=stats["push_gain"],
    )


def run_act(args: argparse.Namespace, outputs_root: Path) -> StageArtifacts:
    from fpgm.objects.interaction import InteractionConfig, InteractionModel, summarize
    from fpgm.robot.urdf import RobotModel

    out_dir = ensure_dir(stage_dir(outputs_root, "act"))
    align_dir = stage_dir(outputs_root, "align")
    _require(align_dir / "alignment.npz", "act", args, "mask,mesh,align")
    alignment = _load_alignment(align_dir)

    traj_path = _load_trajectory_h5(args.episode)
    ctx = _load_clip_context(args.episode, args.camera, args.clip)
    with h5py.File(traj_path, "r") as f:
        joint_positions_full = np.asarray(f["observation/robot_state/joint_positions"])
        gripper_full = np.asarray(f["observation/robot_state/gripper_position"])

    # clip.start/clip.n_frames are in the annotation clock; joint_positions_full is
    # trajectory-clock, so the row indices go through ctx.timing's measured stride
    # (see ClipTiming's docstring) instead of a plain [start, end) slice.
    rows = _strided_rows(
        ctx.clip.start,
        ctx.timing.n_frames,
        ctx.timing.annotation_stride,
        joint_positions_full.shape[0],
        "act stage",
    )
    if rows.size == 0:
        raise DataError(
            f"act stage: clip {args.clip} maps to trajectory rows starting at "
            f"{ctx.clip.start * ctx.timing.annotation_stride}, which is already past "
            f"the end of a {joint_positions_full.shape[0]}-step trajectory"
        )
    joint_positions = joint_positions_full[rows]
    gripper = gripper_full[rows]
    timestamps = ctx.timing.clip_frame_to_seconds(np.arange(rows.size, dtype=np.float64))
    timestamps = timestamps - timestamps[0]  # clip-local seconds, starting at 0

    if not args.urdf.exists():
        raise DataError(
            f"act stage: URDF not found: {args.urdf}. Fetch it with "
            "`python scripts/fetch_robot_description.py --source pointworld`."
        )
    robot = RobotModel(str(args.urdf), load_meshes=False)
    cfg = InteractionConfig(push_gain=args.push_gain)
    trajectory = InteractionModel(robot, cfg).solve(alignment, joint_positions, gripper, timestamps)

    np.savez_compressed(
        out_dir / "trajectory.npz",
        timestamps=trajectory.timestamps, transforms=trajectory.transforms,
        gripper_distance_m=trajectory.gripper_distance_m,
    )
    (out_dir / "states.json").write_text(json.dumps([s.value for s in trajectory.states]))
    summary_text = summarize(trajectory)
    (out_dir / "summary.txt").write_text(summary_text)
    logger.info("interaction timeline:\n%s", summary_text)

    stats = {
        "push_gain": trajectory.push_gain,
        "n_frames": len(trajectory),
        "state_spans": [
            {"state": s.value, "start": a, "end": b} for s, a, b in trajectory.state_spans()
        ],
    }
    (out_dir / "act_stats.json").write_text(json.dumps(stats, indent=2))

    artifacts = StageArtifacts(stage="act", directory=out_dir, stats=stats)
    artifacts.add("trajectory", out_dir / "trajectory.npz")
    artifacts.add("states", out_dir / "states.json")
    artifacts.add("summary_txt", out_dir / "summary.txt")
    artifacts.add("act_stats", out_dir / "act_stats.json")
    return artifacts


# --------------------------------------------------------------------------- #
# Stage 5: render
# --------------------------------------------------------------------------- #


def _object_speeds(trajectory: ObjectTrajectory) -> np.ndarray:
    """Per-frame |velocity| (m/s) of the object's own trajectory, central differences."""
    n = len(trajectory)
    if n < 2:
        return np.zeros(n)
    positions = trajectory.positions
    dt = np.gradient(trajectory.timestamps)
    dt = np.where(dt == 0, np.nan, dt)
    velocity = np.gradient(positions, axis=0) / dt[:, None]
    with np.errstate(invalid="ignore"):
        return np.nan_to_num(np.linalg.norm(velocity, axis=1))


def run_render(args: argparse.Namespace, outputs_root: Path) -> StageArtifacts:
    from fpgm.objects.scene import (
        SceneRenderer,
        draw_point_cloud,
        draw_scene_hud,
        free_camera,
        project_point_cloud,
    )
    from fpgm.robot.overlay import composite
    from fpgm.robot.urdf import RobotModel

    mask_dir = stage_dir(outputs_root, "mask")
    mesh_dir = stage_dir(outputs_root, "mesh")
    align_dir = stage_dir(outputs_root, "align")
    act_dir = stage_dir(outputs_root, "act")
    out_dir = ensure_dir(stage_dir(outputs_root, "render"))

    _require(mask_dir / "meta.json", "render", args, "mask")
    _require(mesh_dir / "mesh.glb", "render", args, "mask,mesh")
    _require(align_dir / "alignment.npz", "render", args, "mask,mesh,align")
    _require(act_dir / "trajectory.npz", "render", args, "mask,mesh,align,act")

    meta = _load_mask_meta(outputs_root)
    mesh = _load_mesh(mesh_dir)
    alignment = _load_alignment(align_dir)
    trajectory = _load_trajectory(act_dir)
    n = len(trajectory)
    if n == 0:
        raise DataError("render stage: act stage produced an empty trajectory")

    if not args.urdf.exists():
        raise DataError(
            f"render stage: URDF not found: {args.urdf}. Fetch it with "
            "`python scripts/fetch_robot_description.py --source pointworld`."
        )
    robot = RobotModel(str(args.urdf), load_meshes=True)

    ctx = _load_clip_context(args.episode, args.camera, args.clip)
    video_w, video_h = meta["video_resolution"]
    camera_video = ctx.camera_annot.rescaled(video_w, video_h)

    traj_path = _load_trajectory_h5(args.episode)
    with h5py.File(traj_path, "r") as f:
        joint_positions_full = np.asarray(f["observation/robot_state/joint_positions"])
        gripper_full = np.asarray(f["observation/robot_state/gripper_position"])
        extrinsics_key = f"observation/camera_extrinsics/{meta['camera_serial']}_left"
        if extrinsics_key not in f:
            raise DataError(f"render stage: trajectory.h5 has no dataset {extrinsics_key!r}")
        extrinsics_6d_full = np.asarray(f[extrinsics_key])

    # `n` frames of the *annotation* clock (act stage's own trajectory length, which
    # may already be shorter than the clip if act stage itself truncated -- see
    # _strided_rows there) map onto trajectory.h5/mp4 rows via ctx.timing's measured
    # stride, not a plain [start, start+n) slice.
    stride = ctx.timing.annotation_stride
    joint_rows = _strided_rows(
        ctx.clip.start, n, stride, joint_positions_full.shape[0], "render stage (joint_positions)"
    )
    extrinsics_rows = _strided_rows(
        ctx.clip.start, n, stride, extrinsics_6d_full.shape[0], "render stage (extrinsics)"
    )
    if joint_rows.size != n or extrinsics_rows.size != n:
        raise DataError(
            f"render stage: trajectory.h5 does not have enough rows to cover all {n} of "
            f"act's frames for clip {args.clip} (mapped rows start at "
            f"{ctx.clip.start * stride}, stride={stride}); re-run act on a shorter clip"
        )
    joint_positions = joint_positions_full[joint_rows]
    gripper = gripper_full[joint_rows]
    extrinsics_6d = extrinsics_6d_full[extrinsics_rows]

    mp4_fps, _count, _size = read_mp4_properties(ctx.mp4_path)
    fps = mp4_fps or _TRAJECTORY_FPS

    target = trajectory.positions.mean(axis=0)
    obj_span = float(np.linalg.norm(mesh.vertices.max(axis=0) - mesh.vertices.min(axis=0)))
    distance = args.freecam_distance or max(0.5, 3.0 * max(obj_span, 0.05))

    speeds = _object_speeds(trajectory)
    point_colors = (
        alignment.point_colors
        if alignment.point_colors is not None
        else np.zeros_like(alignment.points_world, dtype=np.uint8)
    )

    scene_renderer = SceneRenderer(robot.visual_meshes(), mesh, video_w, video_h)
    object_coverage: list[float] = []
    robot_coverage: list[float] = []
    video_start = int(ctx.timing.clip_frame_to_video_frame(0))
    video_end = int(ctx.timing.clip_frame_to_video_frame(n))
    try:
        with ClipFrameSource(str(ctx.mp4_path), video_start, video_end, stride=stride) as frames, \
                VideoWriter(out_dir / "scene_overlay.mp4", fps=fps) as overlay_writer, \
                VideoWriter(out_dir / "scene_freecam.mp4", fps=fps) as freecam_writer:
            if frames.n_frames < n:
                raise DataError(
                    f"render stage: video only yielded {frames.n_frames}/{n} frames for "
                    f"the mapped range [{video_start}, {video_end}) stride={stride}"
                )
            for t in range(n):
                link_poses = robot.link_poses(joint_positions[t], float(gripper[t]))
                object_pose = trajectory.transforms[t]
                state = trajectory.states[t]
                timestamp = float(trajectory.timestamps[t])

                world_to_cam = cam2world_vector_to_world2cam(extrinsics_6d[t])
                camera_t = Camera(camera_video.K, world_to_cam)
                result = scene_renderer.render(link_poses, object_pose, camera_t)
                object_coverage.append(float(result.object_mask.mean()))
                robot_coverage.append(float(result.robot_mask.mean()))

                frame_bgr = frames.read(t)
                composed = composite(frame_bgr, result, alpha=0.85)
                uv, cols = project_point_cloud(
                    alignment.points_world, point_colors, camera_t, (video_h, video_w)
                )
                composed = draw_point_cloud(composed, uv, cols, radius=1)
                composed = draw_scene_hud(
                    composed, t, timestamp, state, float(speeds[t]), trajectory.push_gain
                )
                overlay_writer.write(composed)

                azimuth = 360.0 * t / max(n, 1)
                free_cam = free_camera(
                    target, distance, azimuth, args.freecam_elevation, camera_video.K
                )
                free_result = scene_renderer.render(link_poses, object_pose, free_cam)
                canvas = np.full((video_h, video_w, 3), _FREECAM_BG, dtype=np.uint8)
                canvas = composite(canvas, free_result, alpha=1.0)
                uv_f, cols_f = project_point_cloud(
                    alignment.points_world, point_colors, free_cam, (video_h, video_w)
                )
                canvas = draw_point_cloud(canvas, uv_f, cols_f, radius=1)
                canvas = draw_scene_hud(
                    canvas, t, timestamp, state, float(speeds[t]), trajectory.push_gain
                )
                freecam_writer.write(canvas)
    finally:
        scene_renderer.close()

    stats = {
        "n_frames": n,
        "fps": fps,
        "mean_object_mask_coverage": float(np.mean(object_coverage)) if object_coverage else 0.0,
        "mean_robot_mask_coverage": float(np.mean(robot_coverage)) if robot_coverage else 0.0,
        "peak_object_speed_mps": float(np.max(speeds)) if len(speeds) else 0.0,
        "freecam_distance_m": distance,
    }
    (out_dir / "render_stats.json").write_text(json.dumps(stats, indent=2))

    artifacts = StageArtifacts(stage="render", directory=out_dir, stats=stats)
    artifacts.add("scene_overlay", out_dir / "scene_overlay.mp4")
    artifacts.add("scene_freecam", out_dir / "scene_freecam.mp4")
    artifacts.add("render_stats", out_dir / "render_stats.json")
    return artifacts


# --------------------------------------------------------------------------- #
# Summary + main
# --------------------------------------------------------------------------- #

_RUNNERS = {
    "mask": run_mask,
    "mesh": run_mesh,
    "align": run_align,
    "act": run_act,
    "render": run_render,
}


def _write_summary(outputs_root: Path, args: argparse.Namespace) -> dict:
    """Merge whatever stage artefacts currently exist on disk into one summary.

    Deliberately reads from disk (not from this run's in-memory results), so
    the summary is always the full pipeline's state even when this invocation
    only ran a subset of stages.
    """
    objects_dir = ensure_dir(outputs_root / "objects")
    summary: dict = {"episode": args.episode, "camera": args.camera, "clip": args.clip}

    mask_meta = stage_dir(outputs_root, "mask") / "meta.json"
    if mask_meta.exists():
        summary["mask"] = json.loads(mask_meta.read_text())

    mesh_stats = stage_dir(outputs_root, "mesh") / "mesh_stats.json"
    if mesh_stats.exists():
        summary["mesh"] = json.loads(mesh_stats.read_text())

    align_stats = stage_dir(outputs_root, "align") / "align_stats.json"
    if align_stats.exists():
        summary["align"] = json.loads(align_stats.read_text())

    act_dir = stage_dir(outputs_root, "act")
    act_stats = act_dir / "act_stats.json"
    if act_stats.exists():
        summary["act"] = json.loads(act_stats.read_text())
        summary_txt = act_dir / "summary.txt"
        if summary_txt.exists():
            summary["act"]["timeline"] = summary_txt.read_text()

    render_stats = stage_dir(outputs_root, "render") / "render_stats.json"
    if render_stats.exists():
        summary["render"] = json.loads(render_stats.read_text())

    path = objects_dir / "summary.json"
    path.write_text(json.dumps(summary, indent=2))
    logger.info("summary written to %s", path)
    return summary


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

    stages = [s.strip() for s in args.stages.split(",") if s.strip()]
    unknown = [s for s in stages if s not in STAGES]
    if unknown:
        raise DataError(f"unknown stage(s) {unknown}; choose from {STAGES}")

    outputs_root = args.outputs or (REPO_ROOT / "outputs" / args.episode)
    ensure_dir(outputs_root)
    logger.info(
        "episode=%s camera=%s clip=%s stages=%s", args.episode, args.camera, args.clip, stages
    )

    for stage in stages:
        logger.info("=== stage %d/%d: %s ===", STAGE_INDEX[stage], len(STAGES), stage)
        artifacts = _RUNNERS[stage](args, outputs_root)
        logger.info("%s stage done -> %s", stage, artifacts.directory)

    summary = _write_summary(outputs_root, args)

    print("\n" + "=" * 78)
    print(f"episode:  {args.episode}")
    print(f"camera:   {args.camera}  clip: {args.clip}")
    print(f"stages:   {stages}")
    print(f"outputs:  {outputs_root / 'objects'}")
    if "render" in summary:
        print(json.dumps(summary["render"], indent=2))
    print("=" * 78)
    return 0


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