rpx_benchmark.dataset_hub.ego_layout

  1"""Arrange raw ego (GoPro) captures into the layout scan_capture_root expects.
  2
  3Ego captures land on disk as ``<ego_root>/scene<NNN>/ego/{rgb,sam2}`` — one
  4continuous, non-phased video per scene (see the ego annotation pipeline,
  5``~/Desktop/EGO_ANNOTATION_PIPELINE.md``). That does not match dataset_hub's
  6``<DATA>/{mos,sos}/<scene_id>/<phase>/<modality>`` shape: no top-level
  7``ego/`` scene family recognised before this module, and no numeric phase
  8subdirectory (ego has one continuous clip, not three).
  9
 10This module bridges the two, non-destructively (symlinks by default — nothing
 11here ever writes into the source ego captures):
 12
 13    <DATA>/ego/<scene_id>/0/{rgb,sam2/...}
 14
 15``<scene_id>`` is ego's own bare, zero-padded name (``scene020``) —
 16VERIFIED against the live IRVLUTD/RPX repo, which publishes mos/ scenes as
 17bare ``scene001``..``scene100`` (a local working copy may use a
 18location-suffixed directory name like ``scene20.su.checkerboard.``; the
 19frame/mask bytes underneath are identical to the live repo, only that
 20local directory name differs). Using the bare form means
 21``manifest.py``'s split/difficulty lookup (an exact-string
 22``splits.get(scene.scene_id)``) works against the live splits file with
 23zero extra code — same physical scene, same scene number, captured on
 24independent cameras (rig D435 vs. head-worn GoPro).
 25
 26A few loose files/dirs are intentionally left OUT of the arranged tree —
 27internal to the ego labeling pipeline, not something a downstream benchmark
 28user needs:
 29
 30    ref_frame.txt              which frame the DINOv3 correspondence used
 31                                as its reference (pipeline provenance only)
 32    mask_to_object_multi.json  multi-anchor labeling metadata (only exists
 33                                on a handful of experimentally-relabeled
 34                                scenes; not part of the shipped label set)
 35    jpg/                       secondary JPEG copies of the rgb/ frames
 36
 37``masks_verified/`` has no MOS precedent either — MOS only ships the
 38``verified_masks.txt`` *file* (which both ego and MOS already carry via the
 39existing ``sam2/_meta`` bundle), never a directory of rendered frames.
 40Excluded by default; pass ``include_masks_verified=True`` if that decision
 41changes.
 42"""
 43
 44from __future__ import annotations
 45
 46import re
 47import shutil
 48from dataclasses import dataclass, field
 49from pathlib import Path
 50from typing import List, Optional
 51
 52from ..exceptions import ConfigError
 53
 54#: Loose files directly under sam2/ that are pipeline-internal, not shipped.
 55EXCLUDED_META_FILES = frozenset({"ref_frame.txt", "mask_to_object_multi.json"})
 56
 57#: Directories directly under sam2/ that are pipeline-internal refinement
 58#: provenance (iteration-by-iteration snapshots from manual mask review),
 59#: not final shipped data -- same category as EXCLUDED_META_FILES. Only
 60#: present on scenes that went through >=1 refinement pass; absent on
 61#: clean single-pass scenes. Not recognised by any packer.py modality
 62#: mapping either way, so excluding it is a pure size/time win with no
 63#: behavior change for what actually gets packed.
 64EXCLUDED_META_DIRS = frozenset({"mask_refinement"})
 65
 66#: Top-level dirs under <scene>/ego/ that are never part of the arranged tree
 67#: (only rgb/ and sam2/ are recognised modalities).
 68_ARRANGED_TOP_LEVEL = ("rgb", "sam2")
 69
 70
 71@dataclass(frozen=True)
 72class EgoArrangeResult:
 73    """One scene's outcome from :func:`prepare_ego_layout`."""
 74
 75    ego_scene_dir: str  # e.g. "scene020"
 76    scene_id: Optional[str] = None  # resolved mos/ sibling name, if found
 77    dst: Optional[Path] = None
 78    status: str = "ok"  # "ok" | "no_mos_sibling" | "no_rgb" | "no_sam2"
 79    message: str = ""
 80
 81
 82@dataclass(frozen=True)
 83class EgoLayoutReport:
 84    results: List[EgoArrangeResult] = field(default_factory=list)
 85
 86    @property
 87    def ok(self) -> List[EgoArrangeResult]:
 88        return [r for r in self.results if r.status == "ok"]
 89
 90    @property
 91    def skipped(self) -> List[EgoArrangeResult]:
 92        return [r for r in self.results if r.status != "ok"]
 93
 94
 95_SCENE_NUM_RE = re.compile(r"^scene0*(\d+)$")
 96_MOS_SCENE_NUM_RE = re.compile(r"^scene0*(\d+)(\.|$)")
 97
 98
 99def _ego_scene_number(ego_scene_dir_name: str) -> Optional[int]:
100    m = _SCENE_NUM_RE.match(ego_scene_dir_name)
101    return int(m.group(1)) if m else None
102
103
104def resolve_mos_scene_id(mos_root: Path, ego_scene_dir_name: str) -> Optional[str]:
105    """Validate an ego scene dir has a ``mos/`` sibling, by scene NUMBER, and
106    return the CANONICAL scene_id to use for both.
107
108    IMPORTANT: the canonical scene_id is ego's own bare, zero-padded form
109    (``scene020``) — verified against the live IRVLUTD/RPX repo, where
110    scenes are published as bare ``scene001``..``scene100`` (NOT the
111    location-suffixed form some local working copies use, e.g.
112    ``scene20.su.checkerboard`` — that's a local/staging naming quirk that
113    predates upload; the actual frame/mask bytes underneath match the live
114    repo exactly, only the on-disk directory name differs). Do not use the
115    matched mos/ directory's own (possibly suffixed) name as the scene_id —
116    that was the bug in an earlier version of this function.
117
118    ``mos_root`` is used purely as a validation gate — confirms a scene
119    with this number actually exists as a mos/ scene, regardless of what
120    its local directory happens to be named. Returns ``None`` if there's no
121    exactly-one match (missing, or ambiguous — caller should treat either
122    as "skip, don't guess").
123    """
124    n = _ego_scene_number(ego_scene_dir_name)
125    if n is None or not mos_root.is_dir():
126        return None
127    matches = [
128        d.name for d in mos_root.iterdir()
129        if d.is_dir() and _MOS_SCENE_NUM_RE.match(d.name)
130        and int(_MOS_SCENE_NUM_RE.match(d.name).group(1)) == n
131    ]
132    if len(matches) != 1:
133        return None
134    return ego_scene_dir_name  # canonical bare form, e.g. "scene020"
135
136
137def _link_or_copy(src: Path, dst: Path, mode: str) -> None:
138    dst.parent.mkdir(parents=True, exist_ok=True)
139    if dst.exists() or dst.is_symlink():
140        return
141    if mode == "symlink":
142        if src.is_dir():
143            # Mirror with a REAL directory containing per-file symlinks,
144            # not one directory-level symlink. A directory symlink is
145            # invisible to any standard recursive tree walk that doesn't
146            # explicitly opt in to following symlinks -- notably
147            # ``Path.rglob`` (Python's default: does not descend into
148            # symlinked directories), which lossless_convert.py's
149            # _plan_tree relies on. A single dir-symlink here would make
150            # every rgb/mask frame invisible to lossless-convert while
151            # still resolving fine for direct-path reads (scan/pack) --
152            # confirmed via a real lossless-convert dry-run before this
153            # fix: 0 of ~1500 real ego frames were discovered, only the
154            # two individually-symlinked meta files per scene were. Still
155            # 100% symlinks (nothing written into the source), just at
156            # file granularity so any tree walk works uniformly.
157            dst.mkdir(parents=True, exist_ok=True)
158            for entry in sorted(src.rglob("*")):
159                if not entry.is_file():
160                    continue
161                target = dst / entry.relative_to(src)
162                target.parent.mkdir(parents=True, exist_ok=True)
163                target.symlink_to(entry.resolve())
164        else:
165            dst.symlink_to(src.resolve())
166    elif mode == "copy":
167        if src.is_dir():
168            shutil.copytree(src, dst)
169        else:
170            shutil.copy2(src, dst)
171    else:
172        raise ConfigError(f"unknown mode {mode!r}", hint="Use 'symlink' or 'copy'.")
173
174
175def arrange_ego_scene(
176    ego_scene_src: Path,
177    dst_ego_root: Path,
178    scene_id: str,
179    include_masks_verified: bool = False,
180    mode: str = "symlink",
181) -> Path:
182    """Arrange one ``<ego_root>/scene<NNN>/ego/`` capture into
183    ``<dst_ego_root>/<scene_id>/0/``. Returns the phase-0 dst dir.
184
185    Idempotent: existing links/files at the destination are left alone (so
186    re-running after adding a new scene doesn't touch already-arranged
187    ones). Remove the destination scene dir first if you want a clean
188    re-arrange.
189    """
190    rgb_src = ego_scene_src / "rgb"
191    sam2_src = ego_scene_src / "sam2"
192    if not rgb_src.is_dir():
193        raise ConfigError(f"no rgb/ under {ego_scene_src}")
194    if not sam2_src.is_dir():
195        raise ConfigError(f"no sam2/ under {ego_scene_src}")
196
197    dst_phase = dst_ego_root / scene_id / "0"
198
199    _link_or_copy(rgb_src, dst_phase / "rgb", mode)
200
201    # Allow-list, not exclude-list: ship exactly rgb/ + sam2/masks/ +
202    # mask_to_object.json (+ verified_masks.txt, tiny provenance) --
203    # nothing else. masks_aux (bbox_overlay/dino_output/palette/
204    # contour_gt_masks/rgb_and_mask/masks_contour_with_hidden) is QC/
205    # visualization data, not read by any ego task recipe, and dropped
206    # entirely per the "just rgb + masks (+ the id->name mapping)"
207    # decision -- ego's per-frame composite images are large (high-res
208    # GoPro source), so this is also the single biggest size lever.
209    dst_sam2 = dst_phase / "sam2"
210    masks_src = sam2_src / "masks"
211    if not masks_src.is_dir():
212        raise ConfigError(f"no sam2/masks/ under {ego_scene_src}")
213    _link_or_copy(masks_src, dst_sam2 / "masks", mode)
214    if include_masks_verified and (sam2_src / "masks_verified").is_dir():
215        _link_or_copy(sam2_src / "masks_verified", dst_sam2 / "masks_verified", mode)
216    for fname in ("mask_to_object.json", "verified_masks.txt"):
217        src_file = sam2_src / fname
218        if src_file.is_file():
219            _link_or_copy(src_file, dst_sam2 / fname, mode)
220
221    return dst_phase
222
223
224def prepare_ego_layout(
225    ego_captures_root: Path,
226    mos_root: Path,
227    dst_root: Path,
228    scenes: Optional[List[str]] = None,
229    include_masks_verified: bool = False,
230    mode: str = "symlink",
231) -> EgoLayoutReport:
232    """Arrange every (or a selected subset of) ego scene into
233    ``<dst_root>/ego/<scene_id>/0/`` — ready for ``scan_capture_root``.
234
235    ``scenes``: optional list of ego scene dir names (``["scene020",
236    "scene032"]``) to restrict to. Default: every ``scene*`` dir under
237    ``ego_captures_root``.
238    """
239    ego_captures_root = Path(ego_captures_root)
240    mos_root = Path(mos_root)
241    dst_ego_root = Path(dst_root) / "ego"
242
243    if scenes is None:
244        candidates = sorted(
245            d.name for d in ego_captures_root.iterdir()
246            if d.is_dir() and _ego_scene_number(d.name) is not None
247        )
248    else:
249        candidates = list(scenes)
250
251    results: List[EgoArrangeResult] = []
252    for name in candidates:
253        ego_scene_dir = ego_captures_root / name / "ego"
254        if not (ego_scene_dir / "rgb").is_dir():
255            results.append(EgoArrangeResult(name, status="no_rgb",
256                                             message=f"no rgb/ under {ego_scene_dir}"))
257            continue
258        if not (ego_scene_dir / "sam2" / "masks").is_dir():
259            results.append(EgoArrangeResult(name, status="no_sam2",
260                                             message=f"no sam2/masks/ under {ego_scene_dir}"))
261            continue
262
263        scene_id = resolve_mos_scene_id(mos_root, name)
264        if scene_id is None:
265            results.append(EgoArrangeResult(
266                name, status="no_mos_sibling",
267                message=f"no unique mos/ scene matching {name}'s number",
268            ))
269            continue
270
271        dst = arrange_ego_scene(ego_scene_dir, dst_ego_root, scene_id,
272                                include_masks_verified=include_masks_verified, mode=mode)
273        results.append(EgoArrangeResult(name, scene_id=scene_id, dst=dst))
274
275    return EgoLayoutReport(results=results)