rpx_benchmark.dataset_hub.lossless_convert

  1"""Re-encode a capture tree into denser lossless formats for HF upload.
  2
  3Three modality-aware re-encodings are applied (each can be skipped via
  4:class:`ConvertSpec` toggles). Every encoded artefact is round-trip
  5verified per-frame against the source: the worker decodes the freshly
  6written file and asserts ``numpy.array_equal`` (or full pose-vector
  7equality for cam_pose) before the run is allowed to continue. The first
  8verification failure aborts the entire run with the offending filename
  9and a diagnostic — no silent loss is possible.
 10
 11Modalities and what each becomes
 12--------------------------------
 13
 14* **rgb / fisheye** (and any nested ``rgb/`` like ``ego/rgb/``)
 15  ``*.png`` → ``*.webp`` (libwebp lossless, method=4).
 16  ~45-55% denser on real-world photo content. Decoded pixels are
 17  bit-identical to source via PIL / cv2.
 18
 19* **depth, sam2/masks, sam2/masks_verified**
 20  ``*.png`` → ``*.png`` re-encoded at ``compress_level=9``.
 21  ~18-23% denser; format and mode are preserved exactly — PNG is
 22  lossless at every compression level by spec. Decoded pixel arrays
 23  are bit-identical to source via the same loaders adapters already
 24  use, with no API change.
 25
 26* **cam_pose**
 27  ``*.npz`` (per-frame zip wrapping ``position`` + ``orientation``)
 28  → ``*.npy`` (per-frame raw, single ``(7,) float64`` vector packing
 29  ``[x, y, z, qx, qy, qz, qw]``). ~68% denser per frame (90% of the
 30  current size is per-file zip overhead). Per-frame access is
 31  preserved; the loader dispatches on suffix and remains backward
 32  compatible with the legacy ``.npz`` form.
 33
 34* **everything else** — hard-linked (or copied on cross-filesystem)
 35  into the output tree unchanged. Hard-linking is zero-cost on disk.
 36
 37Why the format zoo is safe
 38--------------------------
 39
 40Every change preserves the decoded pixel/value array the model sees.
 41The encoder asserts ``np.array_equal(src, decoded(encoded(src)))`` at
 42write time, so each artefact on disk is proven bit-faithful before the
 43run continues. PNG-lossless and WebP-lossless are deterministic by
 44spec (RFC 2083 / RFC 9649) — different decoder versions on different
 45machines produce the same pixel arrays. ``.npy`` round-trips IEEE-754
 46floats exactly.
 47
 48Why we did not go further (deferred to later work)
 49--------------------------------------------------
 50
 51The 2-channel hi/lo split that would let depth/masks fit into
 528-bit-WebP-lossless (~+19 GB on the full dataset) was rejected to keep
 53the on-disk file format standard: contributors must not be required to
 54go through the canonical loader to recover correct depth/masks values.
 55That optimisation can be revisited after the loader is wrapped in a
 56typed API that adapters cannot bypass.
 57
 58Usage
 59-----
 60
 61::
 62
 63    from rpx_benchmark.dataset_hub.lossless_convert import (
 64        ConvertSpec, convert_capture_tree,
 65    )
 66    spec = ConvertSpec(src_root=Path("DATA"), out_root=Path("DATA_v2"))
 67    result = convert_capture_tree(spec)
 68    print(f"saved {result.saved_bytes/1e9:.1f} GB")
 69"""
 70
 71from __future__ import annotations
 72
 73import os
 74import shutil
 75from concurrent.futures import ProcessPoolExecutor, as_completed
 76from dataclasses import dataclass, field
 77from pathlib import Path
 78from typing import List, Optional, Tuple
 79
 80from ..exceptions import DatasetError
 81from ..logging_utils import get_logger
 82
 83log = get_logger(__name__)
 84
 85
 86# --------------------------------------------------------------------------- #
 87# Per-file actions
 88# --------------------------------------------------------------------------- #
 89
 90ACTION_WEBP = "webp"  # 8-bit PNG → lossless WebP
 91ACTION_PNG_RECOMPRESS = "png-recompress"  # PNG → PNG at compress_level=9
 92ACTION_CAM_POSE = "cam-pose"  # per-frame .npz → per-frame .npy
 93ACTION_LINK = "link"  # hardlink (or copy)
 94
 95# Modality directory names whose 8-bit PNG frames are re-encoded as
 96# lossless WebP. A file is classified ``webp`` iff its suffix is .png
 97# AND any ancestor directory between src_root and the file is in this
 98# set. Covers ``rgb/``, ``fisheye/``, and nested cases like
 99# ``ego/rgb/`` produced by the GoPro frame extractor.
100WEBP_PARENT_DIRS = frozenset({"rgb", "fisheye"})
101
102# Modality directory names whose PNGs are re-encoded as PNG at
103# compress_level=9 (lossless, mode/dtype preserved, ~20% denser).
104# Depth must stay PNG because libwebp does not support 16-bit;
105# masks must stay PNG to preserve palette/I;16 instance-ID semantics.
106PNG_RECOMPRESS_PARENT_DIRS = frozenset({"depth", "masks", "masks_verified"})
107
108# Directory name that gates the cam_pose path.
109CAM_POSE_PARENT_DIR = "cam_pose"
110
111CONVERT_FROM_SUFFIXES = frozenset({".png", ".PNG"})
112WEBP_SUFFIX = ".webp"
113NPY_SUFFIX = ".npy"
114NPZ_SUFFIX = ".npz"
115
116
117@dataclass(frozen=True)
118class ConvertSpec:
119    """Inputs to :func:`convert_capture_tree`."""
120
121    src_root: Path
122    out_root: Path
123    workers: int = max(1, (os.cpu_count() or 4) - 1)
124    dry_run: bool = False
125    verify: bool = True
126    overwrite_out: bool = False
127    # libwebp lossless effort: 0=fastest/largest, 6=slowest/smallest.
128    # method=4 is the sweet spot on photo content (method=6 takes ~10x
129    # longer for ~0% extra savings on natural images).
130    webp_method: int = 4
131    # PNG re-encode compression level (0-9). 9 = densest, still lossless.
132    png_compress_level: int = 9
133    # Per-modality opt-outs (default: all three paths enabled).
134    skip_rgb_webp: bool = False
135    skip_png_recompress: bool = False
136    skip_cam_pose: bool = False
137
138
139@dataclass
140class ActionStats:
141    """Counts and byte totals for one action."""
142
143    files: int = 0
144    bytes_before: int = 0
145    bytes_after: int = 0
146
147    @property
148    def saved_bytes(self) -> int:
149        return self.bytes_before - self.bytes_after
150
151    @property
152    def saved_pct(self) -> float:
153        if self.bytes_before == 0:
154            return 0.0
155        return 100.0 * self.saved_bytes / self.bytes_before
156
157
158@dataclass
159class ConvertResult:
160    """Outputs of :func:`convert_capture_tree`."""
161
162    files_seen: int = 0
163    by_action: dict = field(default_factory=dict)
164    failures: List[str] = field(default_factory=list)
165
166    def _stats(self, action: str) -> ActionStats:
167        return self.by_action.setdefault(action, ActionStats())
168
169    @property
170    def files_converted(self) -> int:
171        """Files that produced a re-encoded artefact (sum of non-link actions)."""
172        return sum(s.files for a, s in self.by_action.items() if a != ACTION_LINK)
173
174    @property
175    def files_linked(self) -> int:
176        return self._stats(ACTION_LINK).files
177
178    @property
179    def bytes_before(self) -> int:
180        return sum(s.bytes_before for s in self.by_action.values())
181
182    @property
183    def bytes_after(self) -> int:
184        return sum(s.bytes_after for s in self.by_action.values())
185
186    @property
187    def saved_bytes(self) -> int:
188        return self.bytes_before - self.bytes_after
189
190    @property
191    def saved_pct(self) -> float:
192        if self.bytes_before == 0:
193            return 0.0
194        return 100.0 * self.saved_bytes / self.bytes_before
195
196
197# --------------------------------------------------------------------------- #
198# Classification
199# --------------------------------------------------------------------------- #
200
201
202def _classify(path: Path, src_root: Path, spec: Optional[ConvertSpec] = None) -> str:
203    """Return one of the ``ACTION_*`` constants for ``path``.
204
205    Classification is purely structural — it looks at ancestor directory
206    names and file suffix, never at file contents. The optional ``spec``
207    lets a path that *would* be re-encoded fall back to ``ACTION_LINK``
208    when its modality is opted out via a ``skip_*`` toggle.
209    """
210    parts = path.relative_to(src_root).parts[:-1]  # excludes the filename
211
212    # cam_pose: .npz under cam_pose/
213    if path.suffix == NPZ_SUFFIX and CAM_POSE_PARENT_DIR in parts:
214        if spec is not None and spec.skip_cam_pose:
215            return ACTION_LINK
216        return ACTION_CAM_POSE
217
218    # PNG variants
219    if path.suffix in CONVERT_FROM_SUFFIXES:
220        if any(p in WEBP_PARENT_DIRS for p in parts):
221            if spec is not None and spec.skip_rgb_webp:
222                return ACTION_LINK
223            return ACTION_WEBP
224        if any(p in PNG_RECOMPRESS_PARENT_DIRS for p in parts):
225            if spec is not None and spec.skip_png_recompress:
226                return ACTION_LINK
227            return ACTION_PNG_RECOMPRESS
228
229    return ACTION_LINK
230
231
232def _plan_tree(src_root: Path, spec: Optional[ConvertSpec] = None) -> List[Tuple[Path, str]]:
233    """Walk ``src_root`` once and return ``[(path, action), ...]``."""
234    out: List[Tuple[Path, str]] = []
235    for p in sorted(src_root.rglob("*")):
236        if not p.is_file():
237            continue
238        out.append((p, _classify(p, src_root, spec)))
239    return out
240
241
242# --------------------------------------------------------------------------- #
243# Worker functions (top-level so ProcessPoolExecutor can pickle them)
244# --------------------------------------------------------------------------- #
245
246
247def _convert_webp_one(
248    src_path_s: str,
249    dst_path_s: str,
250    verify: bool,
251    webp_method: int,
252) -> Tuple[str, int, int, Optional[str]]:
253    """Encode one PNG → lossless WebP. Optionally round-trip verify.
254
255    Returns ``(dst_str, bytes_before, bytes_after, err_or_None)``.
256    """
257    import numpy as np
258    from PIL import Image
259
260    src_path = Path(src_path_s)
261    dst_path = Path(dst_path_s)
262    try:
263        with Image.open(src_path) as src:
264            src.load()
265            mode = src.mode
266            src_arr = np.asarray(src)
267        before_b = src_path.stat().st_size
268        dst_path.parent.mkdir(parents=True, exist_ok=True)
269        Image.fromarray(src_arr, mode=mode).save(
270            dst_path,
271            format="WEBP",
272            lossless=True,
273            quality=100,
274            method=webp_method,
275        )
276        if verify:
277            with Image.open(dst_path) as restored:
278                restored.load()
279                if restored.mode != mode:
280                    restored = restored.convert(mode)
281                restored_arr = np.asarray(restored)
282            if not np.array_equal(src_arr, restored_arr):
283                return (
284                    dst_path_s,
285                    0,
286                    0,
287                    (
288                        f"WebP round-trip mismatch for {src_path} "
289                        f"(src shape={src_arr.shape} mode={mode}, "
290                        f"restored shape={restored_arr.shape})"
291                    ),
292                )
293        return (dst_path_s, before_b, dst_path.stat().st_size, None)
294    except Exception as e:  # noqa: BLE001 — surface any worker error
295        return (dst_path_s, 0, 0, f"{src_path}: {e!r}")
296
297
298def _recompress_png_one(
299    src_path_s: str,
300    dst_path_s: str,
301    verify: bool,
302    png_compress_level: int,
303) -> Tuple[str, int, int, Optional[str]]:
304    """Re-encode one PNG at higher compress_level. Mode/dtype preserved.
305
306    Returns ``(dst_str, bytes_before, bytes_after, err_or_None)``.
307    """
308    import numpy as np
309    from PIL import Image
310
311    src_path = Path(src_path_s)
312    dst_path = Path(dst_path_s)
313    try:
314        with Image.open(src_path) as src:
315            src.load()
316            src_mode = src.mode
317            src_arr = np.asarray(src)
318            before_b = src_path.stat().st_size
319            dst_path.parent.mkdir(parents=True, exist_ok=True)
320            # Re-save through the original Image so palette / I;16 / L
321            # modes are preserved exactly.
322            src.save(dst_path, format="PNG", compress_level=png_compress_level)
323        if verify:
324            with Image.open(dst_path) as restored:
325                restored.load()
326                restored_mode = restored.mode
327                restored_arr = np.asarray(restored)
328            if restored_mode != src_mode:
329                return (
330                    dst_path_s,
331                    0,
332                    0,
333                    (
334                        f"PNG mode drift for {src_path}: "
335                        f"src mode={src_mode}, restored mode={restored_mode}"
336                    ),
337                )
338            if not np.array_equal(src_arr, restored_arr):
339                return (
340                    dst_path_s,
341                    0,
342                    0,
343                    (
344                        f"PNG round-trip mismatch for {src_path} "
345                        f"(shape={src_arr.shape}, mode={src_mode})"
346                    ),
347                )
348        return (dst_path_s, before_b, dst_path.stat().st_size, None)
349    except Exception as e:  # noqa: BLE001
350        return (dst_path_s, 0, 0, f"{src_path}: {e!r}")
351
352
353def _convert_cam_pose_one(
354    src_path_s: str,
355    dst_path_s: str,
356    verify: bool,
357) -> Tuple[str, int, int, Optional[str]]:
358    """Per-frame ``.npz`` (position + orientation) → ``.npy`` ``(7,) float64``.
359
360    The output is a single float64 vector packing
361    ``[x, y, z, qx, qy, qz, qw]``. Per-frame round-trip verification
362    asserts every slice is bit-identical to the source npz arrays.
363
364    Returns ``(dst_str, bytes_before, bytes_after, err_or_None)``.
365    """
366    import numpy as np
367
368    src_path = Path(src_path_s)
369    dst_path = Path(dst_path_s)
370    try:
371        data = np.load(src_path)
372        if "position" not in data or "orientation" not in data:
373            return (
374                dst_path_s,
375                0,
376                0,
377                (
378                    f"cam_pose .npz missing required keys at {src_path}: "
379                    f"found {list(data.keys())}, need 'position' + 'orientation'"
380                ),
381            )
382        position = data["position"].astype(np.float64)
383        orientation = data["orientation"].astype(np.float64)
384        if position.shape != (3,) or orientation.shape != (4,):
385            return (
386                dst_path_s,
387                0,
388                0,
389                (
390                    f"cam_pose .npz unexpected shapes at {src_path}: "
391                    f"position={position.shape}, orientation={orientation.shape}"
392                ),
393            )
394        pose7 = np.concatenate([position, orientation])  # (7,) float64
395        before_b = src_path.stat().st_size
396        dst_path.parent.mkdir(parents=True, exist_ok=True)
397        np.save(dst_path, pose7)
398        if verify:
399            restored = np.load(dst_path)
400            if restored.shape != (7,) or restored.dtype != np.float64:
401                return (
402                    dst_path_s,
403                    0,
404                    0,
405                    (
406                        f"cam_pose .npy round-trip shape/dtype mismatch at "
407                        f"{dst_path}: shape={restored.shape}, "
408                        f"dtype={restored.dtype}"
409                    ),
410                )
411            if not (
412                np.array_equal(restored[:3], position) and np.array_equal(restored[3:], orientation)
413            ):
414                return (
415                    dst_path_s,
416                    0,
417                    0,
418                    f"cam_pose round-trip value mismatch at {src_path}",
419                )
420        return (dst_path_s, before_b, dst_path.stat().st_size, None)
421    except Exception as e:  # noqa: BLE001
422        return (dst_path_s, 0, 0, f"{src_path}: {e!r}")
423
424
425def _link_one(
426    src_path_s: str,
427    dst_path_s: str,
428) -> Tuple[str, int, int, Optional[str]]:
429    """Hard-link (or copy if cross-fs) a file into the output tree.
430
431    ``src_path`` may itself be a symlink (e.g. ego_layout.py arranges ego
432    captures via per-file symlinks into a scratch DATA tree). On this
433    platform ``os.link(symlink, dst)`` with the default
434    ``follow_symlinks=True`` does NOT dereference -- it creates a new
435    symlink at ``dst`` pointing at the same (possibly relative-to-nowhere)
436    target, not a hardlink to the real file. Confirmed via a direct
437    ``os.lstat`` check: the resulting dst had ``S_ISLNK`` set and the
438    symlink's own ~80-byte size, not the real file's size. That's a
439    silent, severe downstream bug: packer.py's tar writer uses
440    ``os.lstat`` semantics (via ``TarFile.gettarinfo`` without
441    ``dereference=True``), so a symlink dst would get packed as a broken
442    symlink tar entry pointing at a LOCAL machine path -- unreadable by
443    anyone who downloads the resulting shard. Resolve to the real
444    underlying file before linking so this always produces a true
445    hardlink (or a real byte copy, on cross-filesystem fallback) no
446    matter how many symlinks src_path chains through.
447    """
448    src_path = Path(src_path_s).resolve()
449    dst_path = Path(dst_path_s)
450    try:
451        dst_path.parent.mkdir(parents=True, exist_ok=True)
452        if dst_path.exists() or dst_path.is_symlink():
453            dst_path.unlink()
454        try:
455            os.link(src_path, dst_path)
456        except OSError:
457            shutil.copy2(src_path, dst_path)
458        sz = dst_path.stat().st_size
459        return (dst_path_s, sz, sz, None)
460    except Exception as e:  # noqa: BLE001
461        return (dst_path_s, 0, 0, f"{src_path}: {e!r}")
462
463
464# --------------------------------------------------------------------------- #
465# Public entry point
466# --------------------------------------------------------------------------- #
467
468
469def _ensure_disjoint(src: Path, out: Path) -> None:
470    """Refuse src/out overlap — would otherwise lead to infinite walking."""
471    src_r = src.resolve()
472    out_r = out.resolve()
473    if src_r == out_r:
474        raise DatasetError(
475            f"src_root and out_root must differ; got {src_r}",
476            hint="Pass --out pointing at a fresh directory.",
477        )
478    try:
479        if out_r.is_relative_to(src_r) or src_r.is_relative_to(out_r):
480            raise DatasetError(
481                f"src_root and out_root overlap: src={src_r}, out={out_r}",
482                hint="Pick a sibling directory outside src for --out.",
483            )
484    except AttributeError:
485        # Python <3.9 — relative-to check unsupported; fall back to string prefix.
486        if str(out_r).startswith(str(src_r) + os.sep) or str(src_r).startswith(str(out_r) + os.sep):
487            raise DatasetError(
488                f"src_root and out_root overlap: src={src_r}, out={out_r}",
489            ) from None
490
491
492def _dst_path_for(src_path: Path, action: str, src_root: Path, out_root: Path) -> Path:
493    """Return the destination path for ``src_path`` given the action."""
494    rel = src_path.relative_to(src_root)
495    if action == ACTION_WEBP:
496        return out_root / rel.with_suffix(WEBP_SUFFIX)
497    if action == ACTION_CAM_POSE:
498        return out_root / rel.with_suffix(NPY_SUFFIX)
499    # png-recompress and link keep the same filename
500    return out_root / rel
501
502
503def convert_capture_tree(spec: ConvertSpec) -> ConvertResult:
504    """Walk ``spec.src_root`` and produce a re-encoded capture tree.
505
506    On success, ``spec.out_root`` contains a complete capture tree where
507    rgb/fisheye/ego frames are lossless WebP, depth/mask PNGs are
508    re-encoded at maximum compression, and cam_pose frames are
509    consolidated to per-frame .npy. Everything else is hard-linked.
510    The output tree is a drop-in replacement for the source as input to
511    :func:`pack_capture_tree`.
512    """
513    if not spec.src_root.is_dir():
514        raise DatasetError(f"src_root does not exist or is not a directory: {spec.src_root}")
515    _ensure_disjoint(spec.src_root, spec.out_root)
516
517    if spec.out_root.exists() and not spec.overwrite_out:
518        if any(spec.out_root.iterdir()):
519            raise DatasetError(
520                f"out_root exists and is non-empty: {spec.out_root}",
521                hint="Pass overwrite_out=True or pick a fresh directory.",
522            )
523    spec.out_root.mkdir(parents=True, exist_ok=True)
524
525    plan = _plan_tree(spec.src_root, spec=spec)
526    action_counts = {}
527    for _, a in plan:
528        action_counts[a] = action_counts.get(a, 0) + 1
529    log.info(
530        "lossless-convert: scanned %d files under %s; per-action: %s",
531        len(plan),
532        spec.src_root,
533        ", ".join(f"{a}={n}" for a, n in sorted(action_counts.items())),
534    )
535
536    result = ConvertResult(files_seen=len(plan))
537
538    if spec.dry_run:
539        for action, n in action_counts.items():
540            result._stats(action).files = n
541        return result
542
543    # Build per-action worklists.
544    worklists: dict = {
545        ACTION_WEBP: [],
546        ACTION_PNG_RECOMPRESS: [],
547        ACTION_CAM_POSE: [],
548        ACTION_LINK: [],
549    }
550    for src_path, action in plan:
551        dst_path = _dst_path_for(src_path, action, spec.src_root, spec.out_root)
552        worklists[action].append((str(src_path), str(dst_path)))
553
554    # Convert each action's worklist in parallel. Each pass aborts on
555    # first verification failure so a corrupt output tree cannot escape.
556    def run_pass(action: str, fn, *extra_args):
557        worklist = worklists[action]
558        if not worklist:
559            return
560        log.info("lossless-convert: %s — %d files", action, len(worklist))
561        with ProcessPoolExecutor(max_workers=spec.workers) as ex:
562            futures = [ex.submit(fn, src_s, dst_s, *extra_args) for src_s, dst_s in worklist]
563            done = 0
564            for fut in as_completed(futures):
565                _, before_b, after_b, err = fut.result()
566                done += 1
567                if err:
568                    ex.shutdown(wait=False, cancel_futures=True)
569                    raise DatasetError(
570                        f"lossless-convert ({action}) aborted at {done}/{len(worklist)}: {err}",
571                        hint=(
572                            "Inspect the offending file and the relevant codec "
573                            "library version on this machine."
574                        ),
575                    )
576                stats = result._stats(action)
577                stats.files += 1
578                stats.bytes_before += before_b
579                stats.bytes_after += after_b
580                if done % 500 == 0:
581                    log.info("  ... %d / %d done", done, len(worklist))
582
583    run_pass(ACTION_WEBP, _convert_webp_one, spec.verify, spec.webp_method)
584    run_pass(ACTION_PNG_RECOMPRESS, _recompress_png_one, spec.verify, spec.png_compress_level)
585    run_pass(ACTION_CAM_POSE, _convert_cam_pose_one, spec.verify)
586
587    # Linking pass: failures here are collected (not fatal mid-pass) and
588    # raised at the end so we surface every broken file at once.
589    link_failures: List[str] = []
590    if worklists[ACTION_LINK]:
591        log.info("lossless-convert: link — %d files", len(worklists[ACTION_LINK]))
592        with ProcessPoolExecutor(max_workers=spec.workers) as ex:
593            futures = [ex.submit(_link_one, s, d) for s, d in worklists[ACTION_LINK]]
594            for fut in as_completed(futures):
595                _, before_b, after_b, err = fut.result()
596                if err:
597                    link_failures.append(err)
598                    continue
599                stats = result._stats(ACTION_LINK)
600                stats.files += 1
601                stats.bytes_before += before_b
602                stats.bytes_after += after_b
603    if link_failures:
604        result.failures = link_failures
605        raise DatasetError(
606            f"lossless-convert: {len(link_failures)} link/copy failures; first: {link_failures[0]}",
607        )
608
609    return result