rpx_benchmark.pose_pairs

  1"""On-the-fly deterministic pose-pair generation for RPX-RCPE.
  2
  3Generates exact-gap intra-phase pairs and temporal chains without writing
  4anything to disk. Phase 0 and phase 2 are evaluated independently because
  5each capture has its own T265 local world frame; cross-phase relative-pose
  6ground truth is therefore undefined. Given the same seed and config the output is
  7bit-identical, so results are fully reproducible.
  8
  9The generator produces lightweight sample dicts (frame indices + paths)
 10that plug directly into :class:`RPXDataset.from_dict`.  Actual pixel
 11data (RGB, depth) is loaded lazily by the dataset iterator, not here.
 12
 13Usage
 14-----
 15    from rpx_benchmark.pose_pairs import PosePairGenerator, PairConfig
 16
 17    gen = PosePairGenerator(
 18        extracted_root="/path/to/extracted",
 19        parquet_path="/path/to/frames_v1.parquet",
 20        split="easy",
 21    )
 22    manifest = gen.manifest()           # dict consumable by RPXDataset.from_dict
 23    dataset  = gen.as_dataset(batch_size=1)  # ready for BenchmarkRunner
 24
 25    # Or iterate pairs without building the full manifest:
 26    for pair in gen.iter_pairs():
 27        print(pair["id"], pair["pair_type"], pair["rotation_bin"])
 28"""
 29
 30from __future__ import annotations
 31
 32import tarfile
 33from dataclasses import dataclass
 34from pathlib import Path
 35from typing import Any, Dict, Iterator, List, Optional, Set, Tuple
 36
 37import numpy as np
 38import pandas as pd
 39
 40from .exceptions import ConfigError
 41from .logging_utils import get_logger
 42
 43log = get_logger(__name__)
 44
 45
 46# ─────────────────────────────────────────────────────────────────────────────
 47# Tar extraction helpers
 48# ─────────────────────────────────────────────────────────────────────────────
 49
 50def _extract_from_tar(
 51    tar_path: Path, member_name: str, out_path: Path,
 52) -> None:
 53    """Extract a single member from a tar to ``out_path``. Idempotent."""
 54    if out_path.exists():
 55        return
 56    out_path.parent.mkdir(parents=True, exist_ok=True)
 57    with tarfile.open(tar_path, "r") as tf:
 58        f = tf.extractfile(member_name)
 59        if f is None:
 60            raise FileNotFoundError(
 61                f"{member_name} not found in {tar_path}"
 62            )
 63        tmp = out_path.with_suffix(out_path.suffix + ".part")
 64        with tmp.open("wb") as g:
 65            g.write(f.read())
 66        tmp.rename(out_path)
 67
 68RPX_SEED = 5_062_026
 69
 70# ─────────────────────────────────────────────────────────────────────────────
 71# Exclusions
 72# ─────────────────────────────────────────────────────────────────────────────
 73
 74EXCLUDED_SCENE_IDS: Set[str] = {"scene58"}
 75"""Scenes excluded entirely — no ICP-optimized poses available."""
 76
 77EXCLUDED_SCENE_PHASES: Set[Tuple[str, int]] = {
 78    ("scene83.jsom.garden.pot", 0),
 79    ("scene98.jsom.atrium", 0),
 80    ("scene100.jsom.atrium", 0),
 81    ("scene50.ecss.out.stairs", 2),
 82    ("scene72.ecsw.atriumStairs", 2),
 83}
 84"""(scene_id, phase) pairs where ICP optimization failed."""
 85
 86VALID_PHASES: Tuple[int, ...] = (0, 2)
 87"""Clutter and Clean; Interaction is excluded from the RCPE protocol."""
 88
 89
 90def _mode_split(series: pd.Series) -> str | None:
 91    """Return the most common non-null split, or ``None`` if unspecified."""
 92    counts = series.dropna().value_counts()
 93    return str(counts.idxmax()) if not counts.empty else None
 94
 95# ─────────────────────────────────────────────────────────────────────────────
 96# Rotation bins
 97# ─────────────────────────────────────────────────────────────────────────────
 98
 99ROTATION_BINS: List[Tuple[float, float]] = [
100    (0.0, 15.0),
101    (15.0, 45.0),
102    (45.0, 90.0),
103    (90.0, 180.000001),
104]
105BIN_NAMES: List[str] = ["easy", "medium", "hard", "extreme"]
106
107
108def _bin_index(rot_deg: float) -> int:
109    """Return bin index for a rotation, or -1 if below minimum."""
110    for i, (lo, hi) in enumerate(ROTATION_BINS):
111        if lo <= rot_deg < hi:
112            return i
113    return -1
114
115
116# ─────────────────────────────────────────────────────────────────────────────
117# Pose math
118# ─────────────────────────────────────────────────────────────────────────────
119
120def _quat_xyzw_to_rot(q: np.ndarray) -> np.ndarray:
121    x, y, z, w = q / np.linalg.norm(q)
122    return np.array([
123        [1 - 2*(y*y + z*z), 2*(x*y - z*w), 2*(x*z + y*w)],
124        [2*(x*y + z*w), 1 - 2*(x*x + z*z), 2*(y*z - x*w)],
125        [2*(x*z - y*w), 2*(y*z + x*w), 1 - 2*(x*x + y*y)],
126    ], dtype=np.float64)
127
128
129def _load_pose_file(path: Path) -> tuple[np.ndarray, np.ndarray]:
130    """Load an RPX camera pose as ``(rotation, translation)``.
131
132    Published RPX snapshots store compact ``.npy`` vectors as
133    ``[x, y, z, qx, qy, qz, qw]``.  Legacy extracted trees may still contain
134    ``.npz`` files with named ``position`` and ``orientation`` arrays.
135    """
136    data = np.load(path)
137    if path.suffix.lower() == ".npy":
138        values = np.asarray(data, dtype=np.float64)
139        if values.shape != (7,):
140            raise ConfigError(f"camera pose {path} must have shape (7,), got {values.shape}")
141        translation = values[:3]
142        quaternion = values[3:]
143    else:
144        translation = np.asarray(data["position"], dtype=np.float64).reshape(3)
145        quaternion = np.asarray(data["orientation"], dtype=np.float64).reshape(4)
146    return _quat_xyzw_to_rot(quaternion), translation
147
148
149def _rotation_deg(R_a: np.ndarray, R_b: np.ndarray) -> float:
150    cos_t = (np.trace(R_a.T @ R_b) - 1.0) * 0.5
151    return float(np.degrees(np.arccos(np.clip(cos_t, -1.0, 1.0))))
152
153
154# ─────────────────────────────────────────────────────────────────────────────
155# Config
156# ─────────────────────────────────────────────────────────────────────────────
157
158@dataclass
159class PairConfig:
160    """All knobs for pair generation.  Deterministic given ``seed``."""
161
162    # Intra-phase
163    intra_pairs_per_bin: int = 50
164    frame_gap: int = 5
165    # D6 retains every available t -> t+5 pair. The first rotation bin
166    # therefore starts at zero instead of dropping low-motion pairs.
167    min_rotation_deg: float = 0.0
168    min_translation_m: float = 0.0
169
170    # Cross-phase is locked off: separate captures have unrelated T265 worlds.
171    cross_pairs_per_bin: int = 0
172    cross_max_candidates: int = 5000
173
174    # Temporal chains
175    chain_count: int = 5
176    chain_length: int = 10
177    chain_stride: int = 5
178
179    seed: int = RPX_SEED
180
181
182# ─────────────────────────────────────────────────────────────────────────────
183# Core generator
184# ─────────────────────────────────────────────────────────────────────────────
185
186@dataclass
187class _SeqPoses:
188    """Loaded poses + frame indices for one (scene, phase)."""
189    scene: str
190    phase: int
191    frame_idxs: List[int]
192    rotations: List[np.ndarray]   # 3×3 each
193    translations: List[np.ndarray]  # (3,) each
194
195
196class PosePairGenerator:
197    """Deterministic on-the-fly pair generator.
198
199    Loads only poses (NPZ, lightweight) at construction time.
200    Pixel data is never touched — that happens downstream in the
201    :class:`RPXDataset` iterator.
202
203    Parameters
204    ----------
205    extracted_root : Path
206        Root of the extracted frames tree (contains ``scenes/…``).
207    parquet_path : Path
208        ``frames_v1.parquet`` with scene/phase/frame metadata.
209    split : str
210        ``"easy"`` | ``"medium"`` | ``"hard"``.
211    config : PairConfig, optional
212        Generation parameters.  Defaults are the standard benchmark
213        settings (exact t→t+5 pairs plus stride-5 temporal chains).
214    snapshot_root : Path, optional
215        HF snapshot root (parent of ``manifest/`` and tar shards).
216        When provided, cam_pose files are auto-extracted from tars
217        if not already on disk.  If omitted, files must already be
218        extracted under ``extracted_root``.
219    """
220
221    def __init__(
222        self,
223        extracted_root: Path | str,
224        parquet_path: Path | str,
225        split: str,
226        config: PairConfig | None = None,
227        snapshot_root: Path | str | None = None,
228        repo_id: str | None = None,
229    ) -> None:
230        self.root = Path(extracted_root)
231        self.split = split
232        self.cfg = config or PairConfig()
233        if self.cfg.cross_pairs_per_bin != 0:
234            raise ConfigError(
235                "cross-phase RCPE pairs are invalid because each phase has an "
236                "unrelated T265 local world frame; cross_pairs_per_bin must be 0"
237            )
238        self._snapshot_root = Path(snapshot_root) if snapshot_root else None
239        self._repo_id = repo_id
240
241        self._sequences: Dict[Tuple[str, int], _SeqPoses] = {}
242        self._frame_filenames: Dict[Tuple[str, int, int], str] = {}
243        self._cross_scenes: List[str] = []
244        self._load_poses(Path(parquet_path))
245
246    # ── pose loading (one-time) ───────────────────────────────────────────
247
248    def _load_poses(self, parquet_path: Path) -> None:
249        df = pd.read_parquet(parquet_path)
250        if "scene_type" in df.columns:
251            df = df[df["scene_type"].astype(str) == "multi_object"].copy()
252        self._df_cache = df  # kept for ensure_pairs_extracted
253
254        # Scene-wise split assignment
255        scene_split = df.groupby("scene_id")["split"].agg(_mode_split).dropna()
256        df = df.drop(columns=["split"]).merge(
257            scene_split.rename("split"), left_on="scene_id", right_index=True
258        )
259        df = df[
260            (df["split"].astype(str) == self.split)
261            & df["has_cam_pose"].fillna(False).astype(bool)
262        ]
263
264        for (scene, phase), grp in df.groupby(["scene_id", "phase"], sort=True):
265            phase = int(phase)
266            if scene in EXCLUDED_SCENE_IDS:
267                continue
268            if (scene, phase) in EXCLUDED_SCENE_PHASES:
269                continue
270            if phase not in VALID_PHASES:
271                continue
272
273            grp = grp.sort_values("frame_idx").reset_index(drop=True)
274            frame_idxs, rots, trans = [], [], []
275
276            # Auto-extract cam_pose from tar shards if needed
277            self._ensure_cam_pose_extracted(scene, phase, grp)
278
279            for _, row in grp.iterrows():
280                frame_filename = str(row["frame_filename"])
281                stem = frame_filename.rsplit(".", 1)[0]
282                pose_dir = self.root / "scenes" / scene / str(phase) / "cam_pose"
283                npy_path = pose_dir / f"{stem}.npy"
284                npz_path = pose_dir / f"{stem}.npz"
285                pose_path = npy_path if npy_path.exists() else npz_path
286                rotation, translation = _load_pose_file(pose_path)
287                frame_idx = int(row["frame_idx"])
288                frame_idxs.append(frame_idx)
289                rots.append(rotation)
290                trans.append(translation)
291                self._frame_filenames[(str(scene), phase, frame_idx)] = frame_filename
292
293            self._sequences[(scene, phase)] = _SeqPoses(
294                scene=scene, phase=phase,
295                frame_idxs=frame_idxs, rotations=rots, translations=trans,
296            )
297
298        log.info(
299            "loaded poses: %d sequences, split=%s",
300            len(self._sequences), self.split,
301        )
302
303    # ── tar extraction ─────────────────────────────────────────────────────
304
305    def _ensure_cam_pose_extracted(
306        self, scene: str, phase: int, grp: pd.DataFrame,
307    ) -> None:
308        """Extract cam_pose NPZ files from tar shards if not on disk.
309
310        If a tar shard is missing locally (partial HF cache), it is
311        downloaded via ``huggingface_hub`` before extraction.
312        """
313        if self._snapshot_root is None:
314            return
315        if "shard_cam_pose" not in grp.columns:
316            return
317
318        for _, row in grp.iterrows():
319            stem = str(row["frame_filename"]).rsplit(".", 1)[0]
320            pose_dir = (
321                self.root / "scenes" / scene / str(phase)
322                / "cam_pose"
323            )
324            out_path = pose_dir / f"{stem}.npy"
325            if out_path.exists() or (pose_dir / f"{stem}.npz").exists():
326                continue
327            shard_rel = str(row["shard_cam_pose"])
328            tar_path = self._snapshot_root / shard_rel
329            if not tar_path.exists():
330                tar_path = self._download_shard(shard_rel)
331                if tar_path is None:
332                    raise FileNotFoundError(
333                        f"cam_pose shard not found and download failed: {shard_rel}"
334                    )
335            member = f"cam_pose/{stem}.npy"
336            _extract_from_tar(tar_path, member, out_path)
337
338    def _download_shard(self, shard_rel: str) -> Optional[Path]:
339        """Download a single tar shard from HuggingFace Hub.
340
341        Returns the local path on success, None on failure.
342        """
343        try:
344            from huggingface_hub import hf_hub_download
345        except ImportError:
346            log.warning(
347                "huggingface_hub not installed — cannot download shard %s. "
348                "Install with: pip install 'rpx-benchmark[hub]'",
349                shard_rel,
350            )
351            return None
352
353        repo_id = self._repo_id or "IRVLUTD/RPX"
354        log.info("downloading shard %s from %s ...", shard_rel, repo_id)
355        try:
356            local = hf_hub_download(
357                repo_id=repo_id,
358                filename=shard_rel,
359                repo_type="dataset",
360                # Never mix files from the Hub's default branch into a
361                # benchmark pinned to a specific RPX snapshot.
362                revision=(self._snapshot_root.name if self._snapshot_root else None),
363            )
364            return Path(local)
365        except Exception as e:  # noqa: BLE001
366            log.warning("shard download failed for %s: %s", shard_rel, e)
367            return None
368
369    def ensure_pairs_extracted(self, samples: Optional[List[Dict]] = None) -> None:
370        """Extract RGB for all selected pairs so the dataset iterator works.
371
372        Call after :meth:`pairs` or :meth:`manifest`.  Idempotent —
373        skips files already on disk.
374        """
375        if self._snapshot_root is None:
376            return
377        if samples is None:
378            samples = self.pairs()
379
380        df = self._df_cache if hasattr(self, "_df_cache") else None
381        if df is None:
382            return
383
384        # Build a lookup: (scene, phase, stem) → shard_rgb
385        if not hasattr(self, "_shard_lookup"):
386            lookup: Dict[Tuple[str, int, str], str] = {}
387            for _, row in df.iterrows():
388                s = str(row["scene_id"])
389                p = int(row["phase"])
390                stem = str(row["frame_filename"]).rsplit(".", 1)[0]
391                if "shard_rgb" in row.index:
392                    lookup[(s, p, stem)] = str(row["shard_rgb"])
393            self._shard_lookup = lookup
394
395        for s in samples:
396            for path_key in ("rgb", "rgb_b"):
397                rel = s.get(path_key, "")
398                out_path = self.root / rel
399                if out_path.exists():
400                    continue
401                parts = Path(rel).parts
402                if len(parts) < 5:
403                    continue
404                scene, phase_s, _mod, fname = parts[1], parts[2], parts[3], parts[4]
405                stem = fname.rsplit(".", 1)[0]
406                shard = self._shard_lookup.get((scene, int(phase_s), stem))
407                if shard is None:
408                    continue
409                tar_path = self._snapshot_root / shard
410                if not tar_path.exists():
411                    downloaded = self._download_shard(shard)
412                    if downloaded is None:
413                        continue
414                    tar_path = downloaded
415                member = f"rgb/{fname}"
416                _extract_from_tar(tar_path, member, out_path)
417
418    # ── pair generators ───────────────────────────────────────────────────
419
420    def _intra_pairs(
421        self, seq: _SeqPoses, rng: np.random.Generator,
422    ) -> List[Dict[str, Any]]:
423        """Stratified random intra-phase pairs for one sequence."""
424        cfg = self.cfg
425        n = len(seq.frame_idxs)
426        # Collect candidates per bin
427        bins: List[List[Tuple[int, int, float, float]]] = [[] for _ in BIN_NAMES]
428        index_by_frame = {frame: index for index, frame in enumerate(seq.frame_idxs)}
429        for i, frame_a in enumerate(seq.frame_idxs):
430            frame_b = frame_a + cfg.frame_gap
431            j = index_by_frame.get(frame_b)
432            if j is None:
433                continue
434            rot = _rotation_deg(seq.rotations[i], seq.rotations[j])
435            if rot < cfg.min_rotation_deg:
436                continue
437            t_m = float(np.linalg.norm(seq.translations[j] - seq.translations[i]))
438            if t_m < cfg.min_translation_m:
439                continue
440            bi = _bin_index(rot)
441            if bi >= 0:
442                bins[bi].append((frame_a, frame_b, rot, t_m))
443
444        out: List[Dict[str, Any]] = []
445        for bi, name in enumerate(BIN_NAMES):
446            cands = bins[bi]
447            if not cands:
448                continue
449            k = min(len(cands), cfg.intra_pairs_per_bin)
450            sel = rng.choice(len(cands), size=k, replace=False)
451            for idx in sel:
452                fa, fb, rot, t_m = cands[idx]
453                out.append(self._make_entry(
454                    seq.scene, seq.phase, fa, seq.phase, fb,
455                    rot, t_m, "intra_phase", name,
456                ))
457        return out
458
459    def _cross_pairs(
460        self, scene: str, rng: np.random.Generator,
461    ) -> List[Dict[str, Any]]:
462        """Stratified random cross-phase pairs (phase 0 ↔ phase 2)."""
463        cfg = self.cfg
464        s0 = self._sequences[(scene, 0)]
465        s2 = self._sequences[(scene, 2)]
466        n0, n2 = len(s0.frame_idxs), len(s2.frame_idxs)
467        total = n0 * n2
468
469        # Subsample candidate index pairs for speed
470        if total > cfg.cross_max_candidates:
471            flat = rng.choice(total, size=cfg.cross_max_candidates, replace=False)
472            ii, jj = flat // n2, flat % n2
473        else:
474            ii = np.repeat(np.arange(n0), n2)
475            jj = np.tile(np.arange(n2), n0)
476
477        bins: List[List[Tuple[int, int, float, float]]] = [[] for _ in BIN_NAMES]
478        for idx in range(len(ii)):
479            i, j = int(ii[idx]), int(jj[idx])
480            rot = _rotation_deg(s0.rotations[i], s2.rotations[j])
481            if rot < cfg.min_rotation_deg:
482                continue
483            t_m = float(np.linalg.norm(s2.translations[j] - s0.translations[i]))
484            bi = _bin_index(rot)
485            if bi >= 0:
486                bins[bi].append((s0.frame_idxs[i], s2.frame_idxs[j], rot, t_m))
487
488        out: List[Dict[str, Any]] = []
489        for bi, name in enumerate(BIN_NAMES):
490            cands = bins[bi]
491            if not cands:
492                continue
493            k = min(len(cands), cfg.cross_pairs_per_bin)
494            sel = rng.choice(len(cands), size=k, replace=False)
495            for idx in sel:
496                fa, fb, rot, t_m = cands[idx]
497                out.append(self._make_entry(
498                    scene, 0, fa, 2, fb,
499                    rot, t_m, "cross_phase", name,
500                ))
501        return out
502
503    def _temporal_chains(
504        self, seq: _SeqPoses, rng: np.random.Generator,
505    ) -> List[Dict[str, Any]]:
506        """Ordered chains of consecutive pairs for drift evaluation."""
507        cfg = self.cfg
508        index_by_frame = {frame: index for index, frame in enumerate(seq.frame_idxs)}
509        starts = [
510            frame
511            for frame in seq.frame_idxs
512            if all(
513                frame + step * cfg.chain_stride in index_by_frame
514                for step in range(cfg.chain_length + 1)
515            )
516        ]
517        if not starts:
518            return []
519        k = min(cfg.chain_count, len(starts))
520        starts = [starts[index] for index in rng.choice(len(starts), size=k, replace=False)]
521
522        out: List[Dict[str, Any]] = []
523        for ci, start_frame in enumerate(starts):
524            for pi in range(cfg.chain_length):
525                frame_a = start_frame + pi * cfg.chain_stride
526                frame_b = start_frame + (pi + 1) * cfg.chain_stride
527                i = index_by_frame[frame_a]
528                j = index_by_frame[frame_b]
529                rot = _rotation_deg(seq.rotations[i], seq.rotations[j])
530                t_m = float(np.linalg.norm(seq.translations[j] - seq.translations[i]))
531                entry = self._make_entry(
532                    seq.scene, seq.phase, seq.frame_idxs[i],
533                    seq.phase, seq.frame_idxs[j],
534                    rot, t_m, "temporal_chain", None,
535                )
536                # Override ID to include chain index (avoids duplicates
537                # when two chains share a frame pair)
538                sa = f"{seq.frame_idxs[i]:05d}"
539                sb = f"{seq.frame_idxs[j]:05d}"
540                entry["id"] = f"{seq.scene}__{seq.phase}__c{ci}_{pi}__{sa}__{sb}"
541                entry["metadata"]["chain_id"] = (
542                    f"{seq.scene}__{seq.phase}__chain{ci}"
543                )
544                entry["metadata"]["chain_position"] = pi
545                out.append(entry)
546        return out
547
548    # ── entry builder ─────────────────────────────────────────────────────
549
550    def _make_entry(
551        self,
552        scene: str,
553        phase_a: int, frame_a: int,
554        phase_b: int, frame_b: int,
555        rot_deg: float, t_m: float,
556        pair_type: str, rotation_bin: Optional[str],
557    ) -> Dict[str, Any]:
558        frame_filenames = getattr(self, "_frame_filenames", {})
559        filename_a = frame_filenames.get(
560            (str(scene), phase_a, frame_a), f"{frame_a:05d}.png"
561        )
562        filename_b = frame_filenames.get(
563            (str(scene), phase_b, frame_b), f"{frame_b:05d}.png"
564        )
565        sa = Path(filename_a).stem
566        sb = Path(filename_b).stem
567        if pair_type == "cross_phase":
568            pair_id = f"{scene}__cross__0_{sa}__2_{sb}"
569        elif pair_type == "temporal_chain":
570            pair_id = f"{scene}__{phase_a}__tc__{sa}__{sb}"
571        else:
572            pair_id = f"{scene}__{phase_a}__{sa}__{sb}"
573
574        return {
575            "id": pair_id,
576            "scene_id": scene,
577            "phase": phase_a,
578            "phase_b": phase_b,
579            "frame_idx": frame_a,
580            "frame_idx_b": frame_b,
581            "difficulty": None,  # filled by split
582            "pair_type": pair_type,
583            "rotation_bin": rotation_bin,
584            "metadata": {
585                "scene_id": scene,
586                "phase_idx": phase_a,
587                "phase_idx_b": phase_b,
588                "frame": sa,
589                "frame_b": sb,
590                "pair_type": pair_type,
591                "rotation_bin": rotation_bin,
592                "rotation_deg_gt": rot_deg,
593                "translation_m_gt": t_m,
594            },
595            "rgb": f"scenes/{scene}/{phase_a}/rgb/{filename_a}",
596            "rgb_b": f"scenes/{scene}/{phase_b}/rgb/{filename_b}",
597            "pose_a": self._pose_relative_path(scene, phase_a, sa),
598            "pose_b": self._pose_relative_path(scene, phase_b, sb),
599        }
600
601    def _pose_relative_path(self, scene: str, phase: int, stem: str) -> str:
602        """Return the published NPY path, retaining legacy NPZ compatibility."""
603        relative_dir = Path("scenes") / scene / str(phase) / "cam_pose"
604        root = getattr(self, "root", None)
605        if root is not None and (root / relative_dir / f"{stem}.npz").exists():
606            return str(relative_dir / f"{stem}.npz")
607        return str(relative_dir / f"{stem}.npy")
608
609    # ── public API ────────────────────────────────────────────────────────
610
611    def iter_pairs(self) -> Iterator[Dict[str, Any]]:
612        """Yield pair dicts in deterministic order.  Zero disk writes."""
613        rng = np.random.default_rng(self.cfg.seed)
614
615        for key in sorted(self._sequences):
616            seq = self._sequences[key]
617            for p in self._intra_pairs(seq, rng):
618                p["difficulty"] = self.split
619                yield p
620
621        for key in sorted(self._sequences):
622            seq = self._sequences[key]
623            for p in self._temporal_chains(seq, rng):
624                p["difficulty"] = self.split
625                yield p
626
627    def pairs(self) -> List[Dict[str, Any]]:
628        """Materialise all pairs as a list (for manifest / stats)."""
629        return list(self.iter_pairs())
630
631    def manifest(self) -> Dict[str, Any]:
632        """Build a manifest dict consumable by ``RPXDataset.from_dict``.
633
634        No files are written — call ``json.dump`` yourself if you want
635        to cache it, or pass directly to the dataset constructor.
636        """
637        samples = self.pairs()
638        stats = self._compute_stats(samples)
639        return {
640            "task": "relative_camera_pose",
641            "split": self.split,
642            "root": str(self.root),
643            "samples": samples,
644            "_sampler": {
645                "name": "stratified_v2",
646                "version": "2.0",
647                "rotation_bins_deg": [[lo, hi] for lo, hi in ROTATION_BINS],
648                "valid_phases": list(VALID_PHASES),
649                "seed": self.cfg.seed,
650                "intra_pairs_per_bin": self.cfg.intra_pairs_per_bin,
651                "cross_pairs_per_bin": 0,
652                "cross_phase_status": "excluded_unrelated_t265_world_frames",
653                "chain_count": self.cfg.chain_count,
654                "chain_length": self.cfg.chain_length,
655                "chain_stride": self.cfg.chain_stride,
656                "frame_gap": self.cfg.frame_gap,
657                "excluded_scenes": sorted(EXCLUDED_SCENE_IDS),
658                "excluded_scene_phases": [
659                    f"{s}/{p}" for s, p in sorted(EXCLUDED_SCENE_PHASES)
660                ],
661                **stats,
662            },
663        }
664
665    def as_dataset(self, batch_size: int = 1):
666        """Return an :class:`RPXDataset` ready for the runner.
667
668        Automatically extracts RGB files from tar shards for the
669        selected pairs (idempotent, skips existing files).
670        """
671        from .loader import RPXDataset  # noqa: PLC0415
672
673        m = self.manifest()
674        self.ensure_pairs_extracted(m["samples"])
675        return RPXDataset.from_dict(m, batch_size=batch_size)
676
677    @staticmethod
678    def _compute_stats(samples: List[Dict[str, Any]]) -> Dict[str, Any]:
679        by_type = {"intra_phase": 0, "cross_phase": 0, "temporal_chain": 0}
680        by_bin = {n: 0 for n in BIN_NAMES}
681        scenes = set()
682        for s in samples:
683            by_type[s["pair_type"]] += 1
684            if s.get("rotation_bin"):
685                by_bin[s["rotation_bin"]] += 1
686            scenes.add(s["scene_id"])
687        return {
688            "total_pairs": len(samples),
689            "intra_pairs": by_type["intra_phase"],
690            "cross_pairs": by_type["cross_phase"],
691            "temporal_pairs": by_type["temporal_chain"],
692            "pairs_by_bin": by_bin,
693            "n_scenes": len(scenes),
694        }
695
696    def summary(self) -> str:
697        """Human-readable summary string."""
698        m = self.manifest()
699        s = m["_sampler"]
700        lines = [
701            f"RPX-RCPE Pair Generator — split={self.split}",
702            f"  Scenes:           {s['n_scenes']}",
703            f"  Intra-phase:      {s['intra_pairs']:,}",
704            f"  Cross-phase:      {s['cross_pairs']:,}",
705            f"  Temporal chains:  {s['temporal_pairs']:,}",
706            f"  ──────────────────────",
707            f"  TOTAL:            {s['total_pairs']:,}",
708            f"  By bin: {s['pairs_by_bin']}",
709            f"  Seed:             {s['seed']}",
710        ]
711        return "\n".join(lines)