src/fpgm/data/droid_raw.py

"""Client for the public DROID raw GCS bucket.

Downloads per-episode metadata, camera mp4s and ``trajectory.h5`` over plain
anonymous HTTPS (no ``gsutil``, no auth header -- see :mod:`fpgm.data.ids` for
the URL derivation). This module never touches the PointWorld scene-flow
annotations; joining the two sources is :mod:`fpgm.data.episode`'s job.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
from urllib.parse import quote

import numpy as np
from scipy.spatial.transform import Rotation

from fpgm.data.ids import bucket_url_from_scene_path, metadata_filename, resolve_episode_url
from fpgm.types import CameraAssets, DataError, EpisodeId
from fpgm.utils.io import build_retrying_session, download_with_resume
from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

#: Camera roles as named in DROID metadata (``"{role}_cam_serial"`` /
#: ``"{role}_cam_extrinsics"`` / ``"{role}_mp4_path"``).
CAMERA_ROLES = ("wrist", "ext1", "ext2")

#: Which reading of DROID's 6-DOF ``[x, y, z, rx, ry, rz]`` cam-to-world pose to
#: use for the rotation half. **Measured, not assumed**: the matrix was rebuilt
#: under five candidate conventions and compared against PointWorld's stored
#: ground-truth 4x4 for the same camera and episode:
#:
#:     rotvec     max abs err = 1.181
#:     euler_xyz  max abs err = 0.0135   <-- correct
#:     euler_XYZ  max abs err = 1.470
#:
#: The 0.0135 residual is PointWorld's bundle-adjusted refinement of the raw
#: factory calibration, not a convention mismatch. NVlabs/PointWorld's
#: ``real/droid_utils.py`` agrees -- it calls ``convert_pose_euler2mat``.
ROTATION_CONVENTION = "euler_xyz"


def cam2world_vector_to_world2cam(vec: Any) -> np.ndarray:
    """Convert DROID's 6-DOF ``[x, y, z, rx, ry, rz]`` cam-to-world pose to a
    4x4 world-to-camera matrix (standard OpenCV, +Z-forward pinhole).

    DROID stores camera poses as camera-to-world; PointWorld's own convention
    (and the convention every other :class:`~fpgm.types.CameraAssets` in this
    project uses) is world-to-camera, i.e. ``extrinsic = inv(cam2world)``.
    """
    vec = np.asarray(vec, dtype=np.float64).reshape(6)
    translation, rotation = vec[:3], vec[3:6]
    if ROTATION_CONVENTION == "rotvec":
        rot = Rotation.from_rotvec(rotation)
    elif ROTATION_CONVENTION == "euler_xyz":
        rot = Rotation.from_euler("xyz", rotation)
    else:  # pragma: no cover - guarded by the module constant's own two values
        raise ValueError(f"unknown ROTATION_CONVENTION {ROTATION_CONVENTION!r}")

    cam2world = np.eye(4, dtype=np.float64)
    cam2world[:3, :3] = rot.as_matrix()
    cam2world[:3, 3] = translation
    return np.linalg.inv(cam2world)


def read_mp4_fps(path: Path) -> float:
    """Read the container frame rate of an mp4 via OpenCV.

    Must not be assumed equal to the trajectory (control) rate -- see
    :class:`fpgm.types.ClipTiming`'s docstring.
    """
    import cv2

    cap = cv2.VideoCapture(str(path))
    try:
        if not cap.isOpened():
            raise DataError(f"could not open mp4 for fps read: {path}")
        fps = cap.get(cv2.CAP_PROP_FPS)
    finally:
        cap.release()
    if not fps or fps <= 0:
        raise DataError(f"mp4 reports invalid fps ({fps!r}): {path}")
    return float(fps)


def read_mp4_properties(path: Path) -> tuple[float, int, tuple[int, int]]:
    """Read ``(fps, frame_count, (width, height))`` from an mp4 container.

    The frame *count* -- not the fps header -- is what aligns video frames to
    trajectory steps. Real DROID recordings advertise 60 fps while holding one
    frame per 15 Hz control step, so the header alone would stretch the time axis
    by 4x. See :class:`fpgm.types.ClipTiming`.
    """
    import cv2

    cap = cv2.VideoCapture(str(path))
    try:
        if not cap.isOpened():
            raise DataError(f"could not open mp4: {path}")
        fps = float(cap.get(cv2.CAP_PROP_FPS) or 0.0)
        count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
        size = (
            int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
            int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
        )
    finally:
        cap.release()
    return fps, count, size


@dataclass
class RawEpisodeAssets:
    """Locally-downloaded DROID raw assets for one episode.

    Deliberately *not* :class:`fpgm.types.Episode` -- that requires per-camera
    intrinsics and scene-flow clips, which only exist in PointWorld's
    annotations. This is the raw-side half that :class:`fpgm.data.episode.
    EpisodeBuilder` joins against :class:`fpgm.data.pointworld.PointWorldStore`.
    """

    episode_id: EpisodeId
    metadata: dict[str, Any]
    mp4_paths: dict[str, Path] = field(default_factory=dict)  # role -> local path
    trajectory_h5_path: Path | None = None
    camera_serial_by_role: dict[str, str] = field(default_factory=dict)
    # serial -> 4x4 world->camera
    camera_extrinsics: dict[str, np.ndarray] = field(default_factory=dict)


class DroidRawClient:
    """Fetches and locally caches one episode's raw DROID assets."""

    def __init__(
        self,
        cache_dir: Path,
        *,
        timeout: float = 60.0,
        max_retries: int = 5,
        cameras_dir: Path | None = None,
    ):
        """
        Args:
            cameras_dir: PointWorld's ``*_cameras.json`` directory. When given,
                each episode's real bucket path is read from its ``scene_path``
                instead of being derived from the uuid. Measured over all 42935
                episodes, the derivation is wrong for 1.52% (underscores where
                the directory name normally has colons, or a calendar date one
                day off the timestamp) -- those are simply unreachable without
                this, and they concentrate: 14% of IRIS. It also removes the
                success/failure HEAD probe, which 16.4% of episodes fail twice
                over before landing on ``failure/``.
        """
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.timeout = timeout
        self.cameras_dir = Path(cameras_dir) if cameras_dir else None
        self._session = build_retrying_session(max_retries=max_retries)

    def base_url(self, uuid: str) -> str:
        """Bucket base URL for ``uuid``: catalogue first, uuid-derived probe second."""
        if self.cameras_dir is not None:
            cam = self.cameras_dir / f"{uuid}_cameras.json"
            if cam.exists():
                scene_path = json.loads(cam.read_text()).get("scene_path")
                if scene_path:
                    return bucket_url_from_scene_path(scene_path)
        return resolve_episode_url(uuid, timeout=self.timeout)

    def episode_dir(self, uuid: str) -> Path:
        return self.cache_dir / uuid

    def load_episode_metadata(self, uuid: str, *, force_refresh: bool = False) -> dict[str, Any]:
        """Fetch (or read from cache) ``metadata_<uuid>.json``."""
        episode_id = EpisodeId(uuid)
        dest = self.episode_dir(uuid) / "metadata.json"
        if dest.exists() and not force_refresh:
            return json.loads(dest.read_text())

        url = self.base_url(uuid) + metadata_filename(episode_id)
        download_with_resume(url, dest, session=self._session, timeout=self.timeout, progress=False)
        return json.loads(dest.read_text())

    def build_camera_extrinsics(self, metadata: dict[str, Any]) -> dict[str, np.ndarray]:
        """serial -> 4x4 world->camera, converted from DROID's raw metadata poses."""
        extrinsics: dict[str, np.ndarray] = {}
        for role in CAMERA_ROLES:
            serial = metadata.get(f"{role}_cam_serial")
            vec = metadata.get(f"{role}_cam_extrinsics")
            if serial is None or vec is None:
                continue
            extrinsics[serial] = cam2world_vector_to_world2cam(vec)
        return extrinsics

    def download_episode(
        self,
        uuid: str,
        roles: tuple[str, ...] = ("ext1", "ext2"),
        *,
        include_trajectory: bool = True,
    ) -> RawEpisodeAssets:
        """Download metadata, the requested camera mp4s, and (by default) trajectory.h5.

        Files that already exist at their expected remote size are not
        re-downloaded (see :func:`fpgm.utils.io.download_with_resume`).
        """
        episode_id = EpisodeId(uuid)
        metadata = self.load_episode_metadata(uuid)
        base_url = self.base_url(uuid)
        episode_dir = self.episode_dir(uuid)

        camera_serial_by_role: dict[str, str] = {}
        for role in CAMERA_ROLES:
            serial = metadata.get(f"{role}_cam_serial")
            if serial is not None:
                camera_serial_by_role[role] = serial

        mp4_paths: dict[str, Path] = {}
        for role in roles:
            serial = camera_serial_by_role.get(role)
            if serial is None:
                logger.warning("episode %s has no camera for role %r; skipping mp4", uuid, role)
                continue
            mp4_rel = f"recordings/MP4/{serial}.mp4"
            url = base_url + quote(mp4_rel)
            dest = episode_dir / mp4_rel
            download_with_resume(url, dest, session=self._session, timeout=self.timeout)
            mp4_paths[role] = dest

        trajectory_h5_path: Path | None = None
        if include_trajectory:
            url = base_url + "trajectory.h5"
            dest = episode_dir / "trajectory.h5"
            download_with_resume(url, dest, session=self._session, timeout=self.timeout)
            trajectory_h5_path = dest

        return RawEpisodeAssets(
            episode_id=episode_id,
            metadata=metadata,
            mp4_paths=mp4_paths,
            trajectory_h5_path=trajectory_h5_path,
            camera_serial_by_role=camera_serial_by_role,
            camera_extrinsics=self.build_camera_extrinsics(metadata),
        )

    def build_camera_assets(
        self, raw: RawEpisodeAssets, intrinsics: dict[str, np.ndarray]
    ) -> dict[str, CameraAssets]:
        """Assemble :class:`~fpgm.types.CameraAssets` from raw-only sources.

        Raw DROID alone has no intrinsics (those live in PointWorld's
        ``*_flows.h5``), so callers must supply ``intrinsics`` (serial -> 3x3)
        from elsewhere; cameras missing an entry are skipped. Prefer
        :class:`fpgm.data.episode.EpisodeBuilder` over calling this directly --
        it also prefers PointWorld's optimized extrinsics over these raw ones.
        """
        cameras: dict[str, CameraAssets] = {}
        for role, serial in raw.camera_serial_by_role.items():
            intrinsic = intrinsics.get(serial)
            extrinsic = raw.camera_extrinsics.get(serial)
            if intrinsic is None or extrinsic is None:
                continue
            mp4_path = raw.mp4_paths.get(role)
            cameras[serial] = CameraAssets(
                serial=serial,
                role=role,
                intrinsic=intrinsic,
                extrinsic=extrinsic,
                mp4_path=str(mp4_path) if mp4_path is not None else None,
            )
        return cameras