rpx_benchmark.dataset_hub.scanner

  1"""Walk an on-disk RPX capture tree and report structured inventory.
  2
  3Source layout (as captured by the team)::
  4
  5    <root>/
  6    ├── mos/                                  # multi-object scenes
  7    │   └── scene<N>.<building>.<area>/<phase>/<modality>/...
  8    ├── sos/                                  # single-object scenes
  9    │   └── object<N>.<name>/0/<modality>/...
 10    └── ego/                                  # egocentric scenes (one phase, "0")
 11        └── scene<N>.<building>.<area>/0/<modality>/...
 12            # scene_id MUST match its sibling mos/ scene exactly — that's
 13            # how manifest.py joins an ego scene to its MOS split/difficulty
 14            # tier. Use `prepare-ego` (see ego_layout.py) to arrange a raw
 15            # `<scene>/ego/` capture into this shape before scanning.
 16
 17Scene type is determined by which top-level subdirectory the scene lives
 18under (``mos/``, ``sos/``, or ``ego/``) — not by directory name pattern.
 19This makes naming flexible and reduces classification errors.
 20
 21The scanner walks every modality subdir, counts files, sums bytes, and
 22returns a structured ``ScanResult`` that downstream tooling (packer,
 23manifest builder, CLI ``--dry-run``) can format however it likes. The
 24scanner is *purely informational* — it never mutates the tree.
 25"""
 26
 27from __future__ import annotations
 28
 29from dataclasses import dataclass, field
 30from pathlib import Path
 31from typing import Dict, List
 32
 33from .recipes import SceneType
 34
 35# Top-level subdirectory → scene type. Anything else under root (READMEs,
 36# stray captures, half-finished work) is reported via ``ScanResult.skipped``.
 37_TYPE_DIRS: Dict[str, SceneType] = {
 38    "mos": SceneType.MULTI_OBJECT,
 39    "sos": SceneType.SINGLE_OBJECT,
 40    "ego": SceneType.EGO,
 41}
 42
 43#: The inverse of ``_TYPE_DIRS`` — the single source of truth for "where does
 44#: a scene of this type live on disk". Every other module (packer, manifest,
 45#: ...) that needs this mapping should import it from here rather than
 46#: keeping its own copy — a second, un-synced copy is exactly how the
 47#: pre-ego manifest.py bug happened (its own `"mos" if ... else "sos"`
 48#: ternary never learned about a third scene type).
 49SRC_SUBDIR_BY_TYPE: Dict[SceneType, str] = {v: k for k, v in _TYPE_DIRS.items()}
 50
 51
 52@dataclass(frozen=True)
 53class ModalityInventory:
 54    """Inventory of a single modality directory under one phase."""
 55
 56    name: str  # e.g. "rgb", "depth", "fisheye", "cam_pose", "sam2/masks"
 57    file_count: int
 58    total_bytes: int
 59
 60
 61@dataclass(frozen=True)
 62class PhaseInventory:
 63    """Inventory of one phase directory (set of modalities)."""
 64
 65    phase_index: int
 66    modalities: Dict[str, ModalityInventory]
 67
 68    @property
 69    def total_bytes(self) -> int:
 70        return sum(m.total_bytes for m in self.modalities.values())
 71
 72    @property
 73    def file_count(self) -> int:
 74        return sum(m.file_count for m in self.modalities.values())
 75
 76
 77@dataclass(frozen=True)
 78class SceneInventory:
 79    """Inventory of one scene directory (set of phases)."""
 80
 81    scene_id: str
 82    scene_type: SceneType
 83    phases: List[PhaseInventory]
 84
 85    @property
 86    def total_bytes(self) -> int:
 87        return sum(p.total_bytes for p in self.phases)
 88
 89    @property
 90    def file_count(self) -> int:
 91        return sum(p.file_count for p in self.phases)
 92
 93
 94@dataclass(frozen=True)
 95class ScanResult:
 96    """Top-level scan result over a capture root."""
 97
 98    root: Path
 99    scenes: List[SceneInventory]
100    skipped: List[Path] = field(default_factory=list)
101
102    # ── aggregate helpers ────────────────────────────────────────────
103    def by_type(self, scene_type: SceneType) -> List[SceneInventory]:
104        return [s for s in self.scenes if s.scene_type is scene_type]
105
106    @property
107    def total_bytes(self) -> int:
108        return sum(s.total_bytes for s in self.scenes)
109
110    @property
111    def file_count(self) -> int:
112        return sum(s.file_count for s in self.scenes)
113
114    def modality_totals(self) -> Dict[str, ModalityInventory]:
115        """Aggregate per-modality file counts and bytes across all phases."""
116        agg: Dict[str, list[int]] = {}
117        for scene in self.scenes:
118            for phase in scene.phases:
119                for name, inv in phase.modalities.items():
120                    if name not in agg:
121                        agg[name] = [0, 0]
122                    agg[name][0] += inv.file_count
123                    agg[name][1] += inv.total_bytes
124        return {
125            name: ModalityInventory(name=name, file_count=fc, total_bytes=tb)
126            for name, (fc, tb) in sorted(agg.items())
127        }
128
129
130# --------------------------------------------------------------------- #
131# Scanner
132# --------------------------------------------------------------------- #
133
134
135def _measure_dir(d: Path) -> tuple[int, int]:
136    """``(file_count, total_bytes)`` for everything under ``d`` (recursive)."""
137    files, total = 0, 0
138    for p in d.rglob("*"):
139        if p.is_file():
140            files += 1
141            total += p.stat().st_size
142    return files, total
143
144
145def _scan_phase(phase_dir: Path) -> PhaseInventory:
146    """Walk a phase dir and inventory every immediate modality subdir.
147
148    ``sam2/`` is special: each of its subdirs (``masks/``, ``bbox_overlay/``,
149    ...) is reported separately under ``sam2/<sub>`` so the packer can
150    decide which subset to ship per modality. Loose files at the ``sam2/``
151    root (``mask_to_object.json``, ``verified_masks.txt``,
152    ``iter*_faulty.txt``) are bundled into a synthetic ``sam2/_meta``
153    pseudo-modality to keep the inventory complete.
154    """
155    try:
156        idx = int(phase_dir.name)
157    except ValueError:
158        idx = -1
159
160    modalities: Dict[str, ModalityInventory] = {}
161    for child in sorted(phase_dir.iterdir()):
162        if not child.is_dir():
163            continue
164        if child.name == "sam2":
165            for sub in sorted(child.iterdir()):
166                if sub.is_dir():
167                    fc, tb = _measure_dir(sub)
168                    modalities[f"sam2/{sub.name}"] = ModalityInventory(
169                        name=f"sam2/{sub.name}",
170                        file_count=fc,
171                        total_bytes=tb,
172                    )
173            meta_files = [p for p in child.iterdir() if p.is_file()]
174            if meta_files:
175                fc = len(meta_files)
176                tb = sum(p.stat().st_size for p in meta_files)
177                modalities["sam2/_meta"] = ModalityInventory(
178                    name="sam2/_meta",
179                    file_count=fc,
180                    total_bytes=tb,
181                )
182            continue
183        fc, tb = _measure_dir(child)
184        modalities[child.name] = ModalityInventory(
185            name=child.name,
186            file_count=fc,
187            total_bytes=tb,
188        )
189
190    return PhaseInventory(phase_index=idx, modalities=modalities)
191
192
193def scan_capture_root(root: Path) -> ScanResult:
194    """Scan a capture-tree root and return a :class:`ScanResult`.
195
196    Walks ``root/mos/*`` (multi-object scenes) and ``root/sos/*`` (single-
197    object scenes). Anything else under root, and anything inside ``mos/``
198    or ``sos/`` that isn't a scene directory, is collected into
199    :attr:`ScanResult.skipped` so the caller can flag typos / strays.
200    """
201    root = Path(root)
202    if not root.is_dir():
203        raise FileNotFoundError(f"capture root does not exist: {root}")
204
205    scenes: List[SceneInventory] = []
206    skipped: List[Path] = []
207    for entry in sorted(root.iterdir()):
208        # Top-level: only ``mos/`` and ``sos/`` are recognised.
209        if not entry.is_dir() or entry.name not in _TYPE_DIRS:
210            skipped.append(entry)
211            continue
212
213        scene_type = _TYPE_DIRS[entry.name]
214        for scene_dir in sorted(entry.iterdir()):
215            if not scene_dir.is_dir():
216                skipped.append(scene_dir)
217                continue
218            phase_dirs = sorted(
219                [p for p in scene_dir.iterdir() if p.is_dir()],
220                key=lambda p: p.name,
221            )
222            phases = [_scan_phase(p) for p in phase_dirs]
223            scenes.append(
224                SceneInventory(
225                    scene_id=scene_dir.name,
226                    scene_type=scene_type,
227                    phases=phases,
228                )
229            )
230
231    return ScanResult(root=root, scenes=scenes, skipped=skipped)