rpx_benchmark.dataset_hub.mock

  1"""Synthetic mock generator that mimics the on-disk capture layout.
  2
  3Lets us validate the entire dataset-hub pipeline (scan → pack → upload →
  4download → load) without touching the real ~890 GB capture set on the
  5target system. The mock is deliberately tiny (small images, few frames)
  6so a full end-to-end test runs in seconds.
  7
  8Layout produced (matches the ``mos/`` + ``sos/`` wireframe under
  9``benchmark/templates/rpx_capture_wireframe/``)::
 10
 11    <root>/
 12    ├── mos/
 13    │   └── scene1/
 14    │       ├── 0/
 15    │       │   ├── cam_pose/      *.json
 16    │       │   ├── depth/         00000.png ... (16-bit grayscale)
 17    │       │   ├── fisheye/       00000.png ... (8-bit grayscale)
 18    │       │   ├── rgb/           00000.png ... (8-bit RGB)
 19    │       │   └── sam2/
 20    │       │       ├── masks/                       00000.png ...
 21    │       │       ├── bbox_overlay/                00000.png ...
 22    │       │       ├── contour_gt_masks/            00000.png ...
 23    │       │       ├── dino_output/                 00000.png ...
 24    │       │       ├── masks_contour_with_hidden/   00000.png ...
 25    │       │       ├── palette/                     00000.png ...
 26    │       │       ├── rgb_and_mask/                00000.png ...
 27    │       │       ├── mask_to_object.json
 28    │       │       ├── verified_masks.txt
 29    │       │       └── iter1_faulty.txt ... iter4_faulty.txt
 30    │       ├── 1/   (same structure)
 31    │       └── 2/
 32    └── sos/
 33        └── mock_obj_01/                   # bare object name, no "object" prefix
 34            └── 0/
 35                (same modality dirs, one phase only)
 36"""
 37
 38from __future__ import annotations
 39
 40import json
 41import struct
 42import zlib
 43from dataclasses import dataclass
 44from pathlib import Path
 45from typing import List, Optional
 46
 47# --------------------------------------------------------------------- #
 48# Tiny pure-Python PNG writer
 49#
 50# We avoid pulling in PIL/numpy for the mock — the goal is to generate
 51# many small files quickly with zero optional deps. Each function emits
 52# a valid 8-bit-grayscale, 8-bit-RGB, or 16-bit-grayscale PNG.
 53# --------------------------------------------------------------------- #
 54
 55
 56def _png_chunk(tag: bytes, data: bytes) -> bytes:
 57    crc = zlib.crc32(tag + data) & 0xFFFFFFFF
 58    return struct.pack(">I", len(data)) + tag + data + struct.pack(">I", crc)
 59
 60
 61def _png(width: int, height: int, bit_depth: int, color_type: int, raw_rows: List[bytes]) -> bytes:
 62    sig = b"\x89PNG\r\n\x1a\n"
 63    ihdr = struct.pack(">IIBBBBB", width, height, bit_depth, color_type, 0, 0, 0)
 64    raw = b"".join(b"\x00" + row for row in raw_rows)  # filter byte = 0
 65    idat = zlib.compress(raw, 6)
 66    return sig + _png_chunk(b"IHDR", ihdr) + _png_chunk(b"IDAT", idat) + _png_chunk(b"IEND", b"")
 67
 68
 69def _grayscale8_png(width: int, height: int, value: int) -> bytes:
 70    row = bytes([value]) * width
 71    return _png(width, height, 8, 0, [row] * height)
 72
 73
 74def _rgb8_png(width: int, height: int, r: int, g: int, b: int) -> bytes:
 75    row = bytes([r, g, b]) * width
 76    return _png(width, height, 8, 2, [row] * height)
 77
 78
 79def _grayscale16_png(width: int, height: int, value: int) -> bytes:
 80    hi, lo = (value >> 8) & 0xFF, value & 0xFF
 81    row = bytes([hi, lo]) * width
 82    return _png(width, height, 16, 0, [row] * height)
 83
 84
 85# --------------------------------------------------------------------- #
 86# Spec
 87# --------------------------------------------------------------------- #
 88
 89
 90@dataclass(frozen=True)
 91class MockSpec:
 92    """Knobs for the mock generator. Defaults give a ~1 MB tree."""
 93
 94    multi_object_scenes: int = 3
 95    single_object_scenes: int = 5
 96    phases_per_multi: int = 3
 97    frames_per_phase: int = 4
 98    image_size: int = 32
 99    ego_scenes: int = 0
100    """How many of the mos/ scenes also get an ego/ sibling (one phase,
101    "0", same scene_id as its mos/ sibling, same objects). Capped at
102    multi_object_scenes. Default 0 — existing tests are unaffected."""
103    sam2_aux_subdirs: tuple[str, ...] = (
104        "bbox_overlay",
105        "contour_gt_masks",
106        "dino_output",
107        "masks_contour_with_hidden",
108        "palette",
109        "rgb_and_mask",
110    )
111    sam2_iter_files: int = 4
112
113
114# --------------------------------------------------------------------- #
115# Generator
116# --------------------------------------------------------------------- #
117
118
119def generate_mock(out_root: Path, spec: MockSpec | None = None) -> Path:
120    """Materialise a synthetic capture tree under ``out_root``.
121
122    The tree is *deterministic*: re-running with the same spec produces
123    identical bytes (so test assertions on file hashes are stable).
124    """
125    spec = spec or MockSpec()
126    out_root = Path(out_root)
127    out_root.mkdir(parents=True, exist_ok=True)
128    (out_root / "mos").mkdir(exist_ok=True)
129    (out_root / "sos").mkdir(exist_ok=True)
130    if spec.ego_scenes:
131        (out_root / "ego").mkdir(exist_ok=True)
132
133    # SOS first so MOS scenes can reference real SOS object names.
134    sos_object_names: list[str] = []
135    for i in range(1, spec.single_object_scenes + 1):
136        # SOS scene dir name = the bare object name (no "object" prefix).
137        name = f"mock_obj_{i:02d}"
138        sos_object_names.append(name)
139        obj_dir = out_root / "sos" / name
140        _emit_phase(obj_dir / "0", spec, salt=10_000 + i, mask_to_object={"1": name})
141        # FewSOL-style questionnaire at the SOS scene root.
142        (obj_dir / "questionnaire.txt").write_text(
143            _mock_questionnaire(name),
144            encoding="utf-8",
145        )
146
147    # MOS scenes — reference SOS objects in mask_to_object.json so the
148    # questionnaire-dedup layer has something real to look up.
149    for i in range(1, spec.multi_object_scenes + 1):
150        scene_dir = out_root / "mos" / f"scene{i}"
151        # Pick up to 3 of the SOS objects to populate this scene with.
152        if sos_object_names:
153            picked = sos_object_names[
154                (i - 1) % len(sos_object_names) : (i - 1) % len(sos_object_names) + 3
155            ]
156            if not picked:
157                picked = sos_object_names[:3]
158        else:
159            picked = []
160        m2o = {str(idx + 1): name for idx, name in enumerate(picked)}
161        for phase in range(spec.phases_per_multi):
162            _emit_phase(
163                scene_dir / str(phase), spec, salt=i * 100 + phase, mask_to_object=m2o or None
164            )
165
166        # Ego sibling — same scene_id as this mos/ scene (that's the join
167        # key manifest.py uses for split/difficulty), one phase ("0"),
168        # same objects. Salt offset keeps its frames distinct from the
169        # mos phases' pixels.
170        if i <= spec.ego_scenes:
171            _emit_phase(
172                out_root / "ego" / scene_dir.name / "0",
173                spec,
174                salt=i * 100 + 900,
175                mask_to_object=m2o or None,
176            )
177
178    return out_root
179
180
181def _mock_questionnaire(object_name: str) -> str:
182    """Tiny FewSOL-style template so packer/loader tests have real text."""
183    return (
184        f"1. What is the name of the object in these images?\n"
185        f"{object_name}, mock_alias_a, mock_alias_b\n\n"
186        f"2. What is the category of the object in these images?\n"
187        f"mock_category\n\n"
188        f"3. What is the object in these images made of?\n"
189        f"mock_material\n\n"
190        f"4. What can be the object in these images used for?\n"
191        f"mock_use_a, mock_use_b\n\n"
192        f"5. What is the color of the object in these images?\n"
193        f"mock_color_a, mock_color_b\n"
194    )
195
196
197def _emit_phase(
198    phase_dir: Path,
199    spec: MockSpec,
200    salt: int,
201    mask_to_object: Optional[dict[str, str]] = None,
202) -> None:
203    """Write all modality subdirs for one (scene, phase).
204
205    ``mask_to_object`` overrides the synthetic ``sam2/mask_to_object.json``
206    with a caller-supplied mapping. Used by :func:`generate_mock` so MOS
207    scenes reference real SOS object names (the dedup key for the
208    objects_meta/ layer).
209    """
210    rgb_dir = phase_dir / "rgb"
211    rgb_dir.mkdir(parents=True, exist_ok=True)
212    depth_dir = phase_dir / "depth"
213    depth_dir.mkdir(exist_ok=True)
214    # Fisheye is the T265 stereo pair: two subdirs with synced filenames.
215    fisheye_root = phase_dir / "fisheye"
216    fisheye_root.mkdir(exist_ok=True)
217    fisheye_left = fisheye_root / "left"
218    fisheye_left.mkdir(exist_ok=True)
219    fisheye_right = fisheye_root / "right"
220    fisheye_right.mkdir(exist_ok=True)
221    cam_pose_dir = phase_dir / "cam_pose"
222    cam_pose_dir.mkdir(exist_ok=True)
223    sam2_dir = phase_dir / "sam2"
224    sam2_dir.mkdir(exist_ok=True)
225    masks_dir = sam2_dir / "masks"
226    masks_dir.mkdir(exist_ok=True)
227    aux_dirs = {name: (sam2_dir / name) for name in spec.sam2_aux_subdirs}
228    for d in aux_dirs.values():
229        d.mkdir(exist_ok=True)
230
231    n, sz = spec.frames_per_phase, spec.image_size
232    for k in range(n):
233        stem = f"{k:05d}.png"
234        v = (salt + k) & 0xFF
235        (rgb_dir / stem).write_bytes(_rgb8_png(sz, sz, v, (v + 50) & 0xFF, (v + 100) & 0xFF))
236        (depth_dir / stem).write_bytes(_grayscale16_png(sz, sz, ((salt + k) * 37) & 0xFFFF))
237        # Synced filenames: same stem in both fisheye/left/ and fisheye/right/
238        (fisheye_left / stem).write_bytes(_grayscale8_png(sz, sz, (255 - v) & 0xFF))
239        (fisheye_right / stem).write_bytes(_grayscale8_png(sz, sz, v))
240        (masks_dir / stem).write_bytes(_grayscale8_png(sz, sz, (k % 5) + 1))
241        for d in aux_dirs.values():
242            (d / stem).write_bytes(_grayscale8_png(sz, sz, v))
243        (cam_pose_dir / f"{k:05d}.json").write_text(
244            json.dumps(
245                {
246                    "pose": [[1, 0, 0, 0.01 * k], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],
247                    "salt": salt + k,
248                }
249            ),
250            encoding="utf-8",
251        )
252
253    # SAM2 metadata files.
254    m2o = mask_to_object or {str(i): f"mock_object_{i}" for i in range(1, 6)}
255    (sam2_dir / "mask_to_object.json").write_text(
256        json.dumps(m2o, sort_keys=True),
257        encoding="utf-8",
258    )
259    (sam2_dir / "verified_masks.txt").write_text(
260        "\n".join(f"{k:05d}.png" for k in range(n)) + "\n",
261        encoding="utf-8",
262    )
263    for it in range(1, spec.sam2_iter_files + 1):
264        (sam2_dir / f"iter{it}_faulty.txt").write_text(
265            "" if it > 1 else "00000.png\n",
266            encoding="utf-8",
267        )
268
269
270# --------------------------------------------------------------------- #
271# Convenience: total size + file count
272# --------------------------------------------------------------------- #
273
274
275def measure_tree(root: Path) -> tuple[int, int]:
276    """Return ``(file_count, total_bytes)`` under ``root``. Pure stdlib."""
277    files, total = 0, 0
278    for p in Path(root).rglob("*"):
279        if p.is_file():
280            files += 1
281            total += p.stat().st_size
282    return files, total