rpx_benchmark.decode_contracts

  1"""Canonical decode helpers with runtime contracts.
  2
  3The dataset_hub publication pipeline is fault-proof end-to-end —
  4encoder verification, tar SHA-256, per-file SHA-256, manifest-declared
  5extensions. The remaining failure mode is a *consumer-side* one: an
  6adapter that opens files itself, bypasses the loader's contract
  7assertions, and silently feeds wrong-dtype data to a model.
  8
  9The defense is to make the canonical decode path easier to use than
 10the wrong one. Every module that needs a decoded pixel/value array
 11calls one of the ``safe_load_*`` helpers in this module. The helpers
 12encode the *only* correct decode pattern per modality:
 13
 14* :func:`safe_load_rgb`   — ``Image.open(...).convert("RGB")`` →
 15  ``(H, W, 3) uint8``. Any deviation raises ``ManifestError``.
 16* :func:`safe_load_depth` — ``Image.open`` with NO ``.convert()`` so
 17  ``I;16`` mode is preserved → ``(H, W) uint16``. The classic
 18  silent-truncation bug (``cv2.imread`` without ``IMREAD_UNCHANGED``
 19  returning ``uint8``) is impossible inside this helper.
 20* :func:`safe_load_mask`  — palette / ``I;16`` PNG → ``(H, W) int32``
 21  preserving integer instance IDs even past 255.
 22* :func:`safe_load_gray`  — 8-bit grayscale (e.g. fisheye) →
 23  ``(H, W) uint8``. Cannot widen to RGB.
 24* :func:`safe_load_pose`  — ``.npz`` (keys ``position`` + ``orientation``)
 25  or ``.npy`` ``(7,) float64`` packing ``[x, y, z, qx, qy, qz, qw]``.
 26
 27Adapters and per-task scripts that decode files directly *must* go
 28through these helpers. ``RPXDataset._load_*`` in :mod:`loader`
 29delegates to them as well, so there is one and only one place per
 30modality where decode policy lives.
 31
 32If you find yourself reaching for ``cv2.imread`` or ``Image.open``
 33inside an adapter, stop — add a ``safe_load_*`` for the new modality
 34here instead. This module is short by design.
 35"""
 36
 37from __future__ import annotations
 38
 39from pathlib import Path
 40from typing import Tuple
 41
 42import numpy as np
 43from PIL import Image
 44
 45from .exceptions import ManifestError
 46
 47__all__ = [
 48    "safe_load_rgb",
 49    "safe_load_depth",
 50    "safe_load_mask",
 51    "safe_load_gray",
 52    "safe_load_pose",
 53]
 54
 55
 56def safe_load_rgb(path: str | Path) -> np.ndarray:
 57    """Decode an RGB image to ``(H, W, 3) uint8``.
 58
 59    Accepts PNG, WebP-lossless, or any other format Pillow recognises.
 60    The decoded numpy array is bit-identical regardless of container
 61    format, so adapters never need to branch on file suffix.
 62
 63    Raises
 64    ------
 65    ManifestError
 66        If the decoded array is not ``(H, W, 3) uint8``.
 67    """
 68    p = Path(path)
 69    with Image.open(p) as im:
 70        arr = np.array(im.convert("RGB"), dtype=np.uint8)
 71    if arr.dtype != np.uint8 or arr.ndim != 3 or arr.shape[-1] != 3:
 72        raise ManifestError(
 73            f"safe_load_rgb({p}): expected (H, W, 3) uint8, "
 74            f"got shape {arr.shape} dtype {arr.dtype}",
 75        )
 76    return arr
 77
 78
 79def safe_load_depth(path: str | Path) -> np.ndarray:
 80    """Decode a 16-bit single-channel depth PNG to ``(H, W) uint16``.
 81
 82    Does NOT scale or cast to float — that's the caller's job (the
 83    canonical loader divides by 1000 to convert mm → m). This function's
 84    only responsibility is decoding without losing bits.
 85
 86    Critical: never calls ``.convert()``, so PIL's ``I;16`` mode is
 87    preserved. ``cv2.imread`` without ``IMREAD_UNCHANGED`` would
 88    silently return ``uint8`` here; this function refuses that path.
 89
 90    Raises
 91    ------
 92    ManifestError
 93        If the decoded array is not ``(H, W) uint16``.
 94    """
 95    p = Path(path)
 96    with Image.open(p) as im:
 97        arr = np.asarray(im)
 98    if arr.dtype != np.uint16 or arr.ndim != 2:
 99        raise ManifestError(
100            f"safe_load_depth({p}): expected (H, W) uint16 (16-bit PNG), "
101            f"got shape {arr.shape} dtype {arr.dtype}. "
102            f"A uint8 result here would silently truncate every depth "
103            f"value to 8-bit resolution.",
104        )
105    return arr
106
107
108def safe_load_mask(path: str | Path) -> np.ndarray:
109    """Decode an instance-segmentation mask to ``(H, W) int32``.
110
111    Accepts palette PNG (mode ``P``), 16-bit grayscale (``I;16``), or
112    8-bit grayscale (``L``). The int32 cast preserves IDs >255 that
113    palette/I;16 modes can carry.
114
115    Raises
116    ------
117    ManifestError
118        If the decoded array is not ``(H, W) int32``.
119    """
120    p = Path(path)
121    with Image.open(p) as im:
122        arr = np.array(im, dtype=np.int32)
123    if arr.ndim != 2 or arr.dtype != np.int32:
124        raise ManifestError(
125            f"safe_load_mask({p}): expected (H, W) int32, got shape {arr.shape} dtype {arr.dtype}",
126        )
127    return arr
128
129
130def safe_load_gray(path: str | Path) -> np.ndarray:
131    """Decode an 8-bit grayscale image (fisheye) to ``(H, W) uint8``.
132
133    Raises
134    ------
135    ManifestError
136        If the decoded array is not ``(H, W) uint8``.
137    """
138    p = Path(path)
139    with Image.open(p) as im:
140        arr = np.array(im.convert("L"), dtype=np.uint8)
141    if arr.ndim != 2 or arr.dtype != np.uint8:
142        raise ManifestError(
143            f"safe_load_gray({p}): expected (H, W) uint8, got shape {arr.shape} dtype {arr.dtype}",
144        )
145    return arr
146
147
148def safe_load_pose(path: str | Path) -> Tuple[np.ndarray, np.ndarray]:
149    """Decode a T265 pose to ``(position[3], orientation[4])`` float64.
150
151    Accepts both on-disk formats:
152
153    * Legacy ``.npz`` with ``position`` and ``orientation`` keys.
154    * v2 ``.npy`` carrying a ``(7,) float64`` vector packing
155      ``[x, y, z, qx, qy, qz, qw]``.
156
157    Returns
158    -------
159    (position, orientation)
160        ``position`` is ``(3,) float64`` (metres).
161        ``orientation`` is ``(4,) float64`` (xyzw quaternion).
162
163    Raises
164    ------
165    ManifestError
166        If the file's shape/dtype does not match the contract.
167    """
168    p = Path(path)
169    data = np.load(p)
170    if p.suffix.lower() == ".npy":
171        arr = np.asarray(data, dtype=np.float64)
172        if arr.shape != (7,):
173            raise ManifestError(
174                f"safe_load_pose({p}): .npy must be shape (7,) float64; got {arr.shape}",
175            )
176        return arr[:3], arr[3:]
177    # .npz fallback
178    if "position" not in data or "orientation" not in data:
179        raise ManifestError(
180            f"safe_load_pose({p}): .npz missing required keys 'position' and 'orientation'",
181        )
182    position = data["position"].astype(np.float64)
183    orientation = data["orientation"].astype(np.float64)
184    if position.shape != (3,) or orientation.shape != (4,):
185        raise ManifestError(
186            f"safe_load_pose({p}): unexpected shapes "
187            f"position={position.shape}, orientation={orientation.shape}",
188        )
189    return position, orientation