spatialdata-js-util 0.1.0__py3-none-any.whl

spatialdata_js_util/codecs/__init__.py

@@ -0,0 +1,67 @@
+"""Image codecs for browser-oriented SpatialData stores.
+
+`names` holds the on-disk codec identifiers, `backends` picks an HTJ2K
+implementation, `encoding` does the chunk-level work, and `zarr_codec` exposes
+all of it to zarr-python as registered codecs.
+"""
+
+from .backends import (
+    BACKEND_IMAGECODECS,
+    BACKEND_OPENJPH_WASM,
+    backend_report,
+    htj2k_available,
+    resolve_backend,
+)
+from .chunks import chunk_grid, chunk_slices, pad_chunk
+from .encoding import (
+    HTJ2K_DEFAULT_QUALITY_LSB,
+    HTJ2K_QUALITY_FLOOR_LSB,
+    decode_htj2k_plane,
+    decode_image_chunk,
+    decode_image_plane,
+    dtype_quantum,
+    encode_image_chunk,
+    encode_image_plane,
+    htj2k_encode_options,
+)
+from .names import (
+    CODEC_HTJ2K_LEGACY,
+    CODEC_HTJ2K_OPENJPH,
+    CODEC_JPEG2K,
+    HTJ2K_ENCODER,
+    SUPPORTED_IMAGE_CODECS,
+    CodecName,
+    is_htj2k_codec,
+)
+from .zarr_codec import Htj2kCodec, Jpeg2kCodec, LegacyHtj2kCodec, register_codecs
+
+__all__ = [
+    "BACKEND_IMAGECODECS",
+    "BACKEND_OPENJPH_WASM",
+    "CODEC_HTJ2K_LEGACY",
+    "CODEC_HTJ2K_OPENJPH",
+    "CODEC_JPEG2K",
+    "HTJ2K_DEFAULT_QUALITY_LSB",
+    "HTJ2K_ENCODER",
+    "HTJ2K_QUALITY_FLOOR_LSB",
+    "SUPPORTED_IMAGE_CODECS",
+    "Htj2kCodec",
+    "Jpeg2kCodec",
+    "LegacyHtj2kCodec",
+    "CodecName",
+    "backend_report",
+    "chunk_grid",
+    "chunk_slices",
+    "decode_htj2k_plane",
+    "decode_image_chunk",
+    "decode_image_plane",
+    "dtype_quantum",
+    "encode_image_chunk",
+    "encode_image_plane",
+    "htj2k_available",
+    "htj2k_encode_options",
+    "is_htj2k_codec",
+    "pad_chunk",
+    "register_codecs",
+    "resolve_backend",
+]

spatialdata_js_util/codecs/backends.py

@@ -0,0 +1,250 @@
+"""HTJ2K backend selection.
+
+Two backends can produce and consume the HTJ2K codestreams this package writes:
+
+``imagecodecs``
+    Native OpenJPH via `imagecodecs.htj2k_encode/decode`. Pure `pip install` —
+    no Node.js — so it is what makes the reader shim usable from a plain
+    SpatialData environment.
+
+``openjph-wasm``
+    The vendored WASM build driven through a pool of Node.js worker processes.
+    This is the same WASM the JS reader uses, so it is the reference for
+    encode/decode agreement with the browser.
+
+Historically only the WASM backend was trusted, because the WASM decoder this
+repo previously depended on silently dropped every component but the first (see
+`docs/multi-component-codec-findings.md`). That defect is specific to that build,
+not to OpenJPH — but "a build of OpenJPH" is not on its own evidence of
+correctness. So `imagecodecs` is not trusted on reputation: it is admitted only
+if it decodes a committed multi-component codestream to the exact expected
+samples. The probe runs once per process and the backend is rejected on any
+mismatch.
+
+The fixture's expected samples are generated in numpy rather than by a decoder,
+so the ground truth is codec-independent. Which backend *encoded* the codestream
+is recorded alongside it and reported by `backend_report()` — probing a backend
+against a codestream it produced itself only shows self-consistency, so that
+provenance is part of reading the result. See `scripts/build_htj2k_probe.py`.
+
+Set ``SPATIALDATA_JS_UTIL_HTJ2K_BACKEND=imagecodecs|openjph-wasm`` to pin one
+backend explicitly (the probe still gates ``imagecodecs``).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Protocol
+
+import numpy as np
+
+BACKEND_IMAGECODECS = "imagecodecs"
+BACKEND_OPENJPH_WASM = "openjph-wasm"
+
+_BACKEND_ENV_VAR = "SPATIALDATA_JS_UTIL_HTJ2K_BACKEND"
+_PROBE_DIR = Path(__file__).resolve().parent / "probe"
+_PROBE_CODESTREAM = _PROBE_DIR / "multicomponent.j2c"
+_PROBE_EXPECTED = _PROBE_DIR / "multicomponent.npy"
+_PROBE_PROVENANCE = _PROBE_DIR / "multicomponent.json"
+
+
+class Htj2kBackend(Protocol):
+    """Encode/decode HTJ2K codestreams of planar ``(components, y, x)`` volumes."""
+
+    name: str
+
+    def available(self) -> bool: ...
+
+    def encode(self, volume: np.ndarray, *, reversible: bool, quality: float) -> bytes: ...
+
+    def decode(self, codestream: bytes | bytearray) -> np.ndarray: ...
+
+
+def _as_planar_3d(array: np.ndarray) -> np.ndarray:
+    """Normalise a decoded result to ``(components, y, x)``."""
+    if array.ndim == 2:
+        return array.reshape(1, *array.shape)
+    if array.ndim == 3:
+        return array
+    raise ValueError(f"Expected a 2D or 3D HTJ2K decode result, got shape {array.shape}")
+
+
+class ImagecodecsHtj2kBackend:
+    """Native OpenJPH through `imagecodecs`; no Node.js required."""
+
+    name = BACKEND_IMAGECODECS
+
+    def available(self) -> bool:
+        try:
+            import imagecodecs
+        except ImportError:
+            return False
+        return bool(getattr(imagecodecs, "HTJ2K", None) and imagecodecs.HTJ2K.available)
+
+    def encode(self, volume: np.ndarray, *, reversible: bool, quality: float) -> bytes:
+        import imagecodecs
+
+        array = np.ascontiguousarray(np.asarray(volume))
+        planar = array.ndim == 3 and array.shape[0] > 1
+        if array.ndim == 3 and array.shape[0] == 1:
+            array = array.reshape(array.shape[1], array.shape[2])
+        if reversible:
+            # `level` is the OpenJPH quantization step; omitting it selects the
+            # reversible (5/3) transform.
+            return bytes(imagecodecs.htj2k_encode(array, reversible=True, planar=planar, rgb=False))
+        return bytes(
+            imagecodecs.htj2k_encode(
+                array, level=float(quality), reversible=False, planar=planar, rgb=False
+            )
+        )
+
+    def decode(self, codestream: bytes | bytearray) -> np.ndarray:
+        import imagecodecs
+
+        return _as_planar_3d(np.asarray(imagecodecs.htj2k_decode(bytes(codestream), planar=True)))
+
+
+class OpenJphWasmHtj2kBackend:
+    """Vendored openjph-wasm driven by Node.js worker processes."""
+
+    name = BACKEND_OPENJPH_WASM
+
+    def available(self) -> bool:
+        from .htj2k_wasm import wasm_encode_available
+
+        return wasm_encode_available()
+
+    def encode(self, volume: np.ndarray, *, reversible: bool, quality: float) -> bytes:
+        from .htj2k_wasm import encode_htj2k_wasm
+
+        return encode_htj2k_wasm(volume, reversible=reversible, quality=quality)
+
+    def decode(self, codestream: bytes | bytearray) -> np.ndarray:
+        from .htj2k_wasm import decode_htj2k_wasm
+
+        return _as_planar_3d(decode_htj2k_wasm(codestream))
+
+
+_BACKENDS: dict[str, Htj2kBackend] = {
+    BACKEND_IMAGECODECS: ImagecodecsHtj2kBackend(),
+    BACKEND_OPENJPH_WASM: OpenJphWasmHtj2kBackend(),
+}
+
+# `imagecodecs` first so that a plain `pip install` can both read and write
+# without Node.js. The probe below is what makes that safe.
+_PREFERENCE = (BACKEND_IMAGECODECS, BACKEND_OPENJPH_WASM)
+
+
+def probe_fixture() -> tuple[bytes, np.ndarray]:
+    """Return the committed multi-component codestream and its expected samples."""
+    if not _PROBE_CODESTREAM.is_file() or not _PROBE_EXPECTED.is_file():
+        raise FileNotFoundError(
+            f"HTJ2K probe fixture missing under {_PROBE_DIR}. "
+            "Regenerate it with scripts/build_htj2k_probe.py."
+        )
+    return _PROBE_CODESTREAM.read_bytes(), np.load(_PROBE_EXPECTED)
+
+
+def probe_encoder() -> str | None:
+    """Which backend encoded the committed fixture, or `None` if unrecorded.
+
+    A backend probed against a codestream it produced itself has only shown that
+    it is self-consistent, so this is needed to read a probe result correctly.
+    """
+    if not _PROBE_PROVENANCE.is_file():
+        return None
+    try:
+        encoder = json.loads(_PROBE_PROVENANCE.read_text()).get("encoder")
+    except (OSError, ValueError):
+        return None
+    return encoder if isinstance(encoder, str) else None
+
+
+def available_backends() -> list[Htj2kBackend]:
+    """Every installed backend, in preference order.
+
+    Exposed for the probe builder, which has to encode with one backend and
+    check the result against the others rather than take the single selected one.
+    """
+    return [_BACKENDS[name] for name in _PREFERENCE if _BACKENDS[name].available()]
+
+
+@lru_cache(maxsize=None)
+def backend_passes_probe(name: str) -> bool:
+    """Whether *name* decodes the committed multi-component fixture exactly.
+
+    A backend that raises, returns the wrong shape, or returns wrong samples
+    fails. Multi-component correctness cannot be assumed from the library name:
+    the WASM build previously used here decoded every component as component 0.
+    """
+    backend = _BACKENDS.get(name)
+    if backend is None or not backend.available():
+        return False
+    try:
+        codestream, expected = probe_fixture()
+        decoded = backend.decode(codestream)
+    except Exception:
+        return False
+    return decoded.shape == expected.shape and np.array_equal(decoded, expected)
+
+
+@lru_cache(maxsize=1)
+def resolve_backend() -> Htj2kBackend | None:
+    """Return the HTJ2K backend to use, or ``None`` if none is usable."""
+    requested = os.environ.get(_BACKEND_ENV_VAR)
+    candidates = (requested,) if requested else _PREFERENCE
+    if requested and requested not in _BACKENDS:
+        raise ValueError(
+            f"Unknown {_BACKEND_ENV_VAR}={requested!r}; expected one of {sorted(_BACKENDS)}"
+        )
+    for name in candidates:
+        # The WASM backend is the reference the probe fixture was produced with,
+        # so it is trusted on availability; every other backend must prove it.
+        if name == BACKEND_OPENJPH_WASM:
+            if _BACKENDS[name].available():
+                return _BACKENDS[name]
+        elif backend_passes_probe(name):
+            return _BACKENDS[name]
+    return None
+
+
+def reset_backend_cache() -> None:
+    """Clear memoized backend selection (used by tests that flip the env var)."""
+    resolve_backend.cache_clear()
+    backend_passes_probe.cache_clear()
+
+
+def require_backend() -> Htj2kBackend:
+    backend = resolve_backend()
+    if backend is None:
+        raise RuntimeError(
+            "No usable HTJ2K backend. Install `imagecodecs` with HTJ2K support "
+            "(pip install 'spatialdata-js-util'), or put Node.js on PATH to use the "
+            "vendored openjph-wasm backend."
+        )
+    return backend
+
+
+def htj2k_available() -> bool:
+    return resolve_backend() is not None
+
+
+def backend_report() -> dict[str, Any]:
+    """Describe backend availability — used by manifests and `codecs info`."""
+    selected = resolve_backend()
+    return {
+        "selected": selected.name if selected else None,
+        # Reading a probe result needs to know who produced the codestream: a
+        # backend that encoded it has only proved self-consistency by decoding it.
+        "probe_encoder": probe_encoder(),
+        "backends": {
+            name: {
+                "available": backend.available(),
+                "passes_multicomponent_probe": backend_passes_probe(name),
+            }
+            for name, backend in _BACKENDS.items()
+        },
+    }

spatialdata_js_util/codecs/chunks.py

@@ -0,0 +1,40 @@
+"""Chunk-grid arithmetic shared by the recompressor and its tests."""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def chunk_grid(shape: tuple[int, ...], chunks: tuple[int, ...]) -> list[tuple[int, ...]]:
+    if len(shape) != len(chunks):
+        raise ValueError(
+            f"shape and chunks must have the same length (got {len(shape)} and {len(chunks)})"
+        )
+    ranges = [range((size + chunk - 1) // chunk) for size, chunk in zip(shape, chunks)]
+    out: list[tuple[int, ...]] = [()]
+    for values in ranges:
+        out = [(*prefix, value) for prefix in out for value in values]
+    return out
+
+
+def chunk_slices(
+    shape: tuple[int, ...], chunks: tuple[int, ...], coords: tuple[int, ...]
+) -> tuple[slice, ...]:
+    if not (len(shape) == len(chunks) == len(coords)):
+        raise ValueError(
+            "shape, chunks, and coords must have the same length "
+            f"(got {len(shape)}, {len(chunks)}, and {len(coords)})"
+        )
+    slices = []
+    for coord, chunk, size in zip(coords, chunks, shape):
+        start = coord * chunk
+        slices.append(slice(start, min(start + chunk, size)))
+    return tuple(slices)
+
+
+def pad_chunk(chunk: np.ndarray, chunks: tuple[int, ...]) -> np.ndarray:
+    if chunk.shape == chunks:
+        return chunk
+    padded = np.zeros(chunks, dtype=chunk.dtype)
+    padded[tuple(slice(0, size) for size in chunk.shape)] = chunk
+    return padded

spatialdata_js_util/codecs/encoding.py

@@ -0,0 +1,154 @@
+"""Encode and decode image chunks for the codecs this package supports.
+
+JPEG 2000 goes through `imagecodecs`. HTJ2K goes through whichever backend
+`backends.resolve_backend()` selects. Both directions are expressed in terms of
+planar ``(components, y, x)`` volumes; 2D planes are the single-component case.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from .backends import require_backend
+from .names import CODEC_HTJ2K_OPENJPH, CODEC_JPEG2K, is_htj2k_codec
+
+
+# A quantization step below one input LSB cannot buy fidelity the input does not
+# have: the irreversible 9/7 path returns a bit-identical image while spending
+# more bytes than the reversible 5/3 path. See `images.HTJ2K_PRESETS`.
+HTJ2K_QUALITY_FLOOR_LSB = 1.0
+
+# Fallback for `reversible=False` with no step given, in LSB. Matches the
+# `balanced` preset.
+HTJ2K_DEFAULT_QUALITY_LSB = 2.0
+
+
+def dtype_quantum(dtype: np.dtype) -> float:
+    """Return one LSB of *dtype* as a fraction of its full dynamic range.
+
+    OpenJPH's quantization step is normalised to the dtype's full range, so this
+    is the unit that makes a step comparable across bit depths.
+
+    Only integer dtypes have an LSB to be relative to. That reflects what this
+    codec path writes today, not a limit of HTJ2K — if float support is added,
+    give those dtypes a quantum here rather than reading the error as permanent.
+    """
+    resolved = np.dtype(dtype)
+    if resolved.kind not in ("i", "u"):
+        raise TypeError(
+            f"HTJ2K quantization steps are relative to an integer dtype's full range, "
+            f"and {resolved} has no LSB to scale by. Images are currently written as "
+            f"<=16-bit integers, so no step is derived for this dtype yet — pass an "
+            f"explicit `quality` (an absolute step) or `reversible=True` to encode it."
+        )
+    info = np.iinfo(resolved)
+    return 1.0 / (int(info.max) - int(info.min) + 1)
+
+
+def htj2k_encode_options(
+    encode_options: dict[str, Any], *, dtype: np.dtype | None = None
+) -> tuple[bool, float]:
+    """Map encode options to the OpenJPH ``(reversible, quality)`` pair.
+
+    ``quality`` is the OpenJPH quantization step: lower means higher fidelity and
+    a larger codestream. It is *not* a JPEG-style 0–100 quality.
+
+    A caller-supplied ``quality``/``level`` is used as given. Only the fallback
+    is derived from *dtype*, so that the default scales with bit depth instead of
+    being a bit-depth-blind constant; deriving it needs an integer dtype.
+    """
+    reversible = bool(encode_options.get("reversible", True))
+    if reversible:
+        return True, 0.0
+    quality = encode_options.get("quality", encode_options.get("level"))
+    if quality is None:
+        # Derived only when the caller gave no step. An explicit `quality` is an
+        # absolute value, so it must not depend on the dtype — deriving eagerly
+        # would reject dtypes that can be encoded perfectly well by naming a step.
+        quality = HTJ2K_DEFAULT_QUALITY_LSB * dtype_quantum(dtype) if dtype is not None else 0.0002
+    return False, float(quality)
+
+
+def _encode_htj2k(volume: np.ndarray, encode_options: dict[str, Any] | None = None) -> bytes:
+    reversible, quality = htj2k_encode_options(encode_options or {}, dtype=volume.dtype)
+    return require_backend().encode(volume, reversible=reversible, quality=quality)
+
+
+def _htj2k_components(array: np.ndarray) -> int:
+    return int(np.prod(array.shape[:-2])) if array.ndim > 2 else 1
+
+
+def decode_htj2k_plane(encoded: bytes | bytearray) -> np.ndarray:
+    """Decode an HTJ2K codestream.
+
+    Returns a 2D ``(y, x)`` array for single-component codestreams, otherwise a
+    ``(components, y, x)`` array.
+    """
+    array = require_backend().decode(encoded)
+    return array[0] if array.shape[0] == 1 else array
+
+
+def decode_image_plane(encoded: bytes | bytearray, codec: str) -> np.ndarray:
+    if codec == CODEC_JPEG2K:
+        import imagecodecs
+
+        return imagecodecs.jpeg2k_decode(encoded)
+    if is_htj2k_codec(codec):
+        return decode_htj2k_plane(encoded)
+    raise ValueError(f"Unsupported image codec: {codec}")
+
+
+def encode_image_plane(
+    plane: np.ndarray, codec: str, encode_options: dict[str, Any]
+) -> bytes | bytearray:
+    array = np.asarray(plane)
+    if codec == CODEC_JPEG2K:
+        import imagecodecs
+
+        return imagecodecs.jpeg2k_encode(array, **encode_options)
+    if codec == CODEC_HTJ2K_OPENJPH:
+        return _encode_htj2k(array, encode_options)
+    raise ValueError(f"Unsupported image codec: {codec}")
+
+
+def encode_image_chunk(
+    volume: np.ndarray, codec: str, encode_options: dict[str, Any]
+) -> bytes | bytearray:
+    """Encode one or more planar components to a single codestream.
+
+    ``volume`` is ``(components, y, x)`` (or 2D for a single component). HTJ2K
+    encodes the planes as codestream components (e.g. z-planes of a volumetric
+    chunk); multi-component JPEG2K chunks are not produced by this writer.
+    """
+    array = np.ascontiguousarray(np.asarray(volume))
+    components = _htj2k_components(array)
+    if codec == CODEC_HTJ2K_OPENJPH:
+        return _encode_htj2k(array, encode_options)
+    if codec == CODEC_JPEG2K:
+        import imagecodecs
+
+        if components == 1:
+            plane = array.reshape(array.shape[-2], array.shape[-1])
+            return imagecodecs.jpeg2k_encode(plane, **encode_options)
+        raise NotImplementedError(
+            "Multi-component JPEG2K chunks are not supported; use HTJ2K (openjph)."
+        )
+    raise ValueError(f"Unsupported image codec: {codec}")
+
+
+def decode_image_chunk(
+    encoded: bytes | bytearray, codec: str, *, components: int, height: int, width: int
+) -> np.ndarray:
+    """Decode a chunk codestream to a ``(components, y, x)`` array."""
+    if is_htj2k_codec(codec):
+        return require_backend().decode(encoded).reshape(components, height, width)
+    if codec == CODEC_JPEG2K:
+        import imagecodecs
+
+        decoded = np.asarray(imagecodecs.jpeg2k_decode(encoded))
+        if decoded.ndim == 2:
+            return decoded.reshape(1, height, width)
+        return np.moveaxis(decoded, -1, 0).reshape(components, height, width)
+    raise ValueError(f"Unsupported image codec: {codec}")