Add zarr codec support for JP2K/HTJ2K with Python writer implementation

.changeset/legal-moles-design.md

@@ -0,0 +1,10 @@
+---
+"@spatialdata/avivatorish": patch
+"zarrextra": patch
+"@spatialdata/core": patch
+"@spatialdata/vis": patch
+---
+
+Add support for alternative codecs in zarrextra, with tooling to encode images as JPEG2000 and HTJ2K.
+
+Zarrita stores can be configured to decode in workers.

.cursor/settings.json

@@ -9,7 +9,8 @@
     "editor.formatOnSave": true
   },
   "python.analysis.extraPaths": [
-    "${workspaceFolder}/python"
+    "${workspaceFolder}/python",
+    "${workspaceFolder}/python/spatialdata-codec-writer/scripts"
   ]
 }
 

.github/workflows/test.yml

@@ -35,6 +35,9 @@ jobs:
       - name: Install dependencies
         run: pnpm install --frozen-lockfile
 
+      - name: Vendor OpenJPH for spatialdata-codec-writer
+        run: node scripts/vendor-openjph-for-python.mjs
+
       - name: Build packages
         run: pnpm build
 
@@ -51,6 +54,9 @@ jobs:
       - name: Generate test fixtures (Python spatialdata 0.5.0, 0.6.1, and 0.7.2)
         run: pnpm test:fixtures:generate
 
+      - name: Run spatialdata-codec-writer tests
+        run: uv run --directory python/spatialdata-codec-writer pytest
+
       - name: Run unit tests
         run: pnpm test:unit
 

.gitignore

@@ -19,6 +19,13 @@ python/v0.5.0/.venv/
 python/v0.6.1/.venv/
 python/v0.7.2/.venv/
 
+# Vendored from @cornerstonejs/codec-openjph (see scripts/vendor-openjph-for-python.mjs)
+python/spatialdata-codec-writer/src/spatialdata_codec_writer/vendor/openjph/
+
+**/__pycache__/
+*.egg-info/
+*.egg
+
 # Environment
 .env
 .env.local

.vscode/settings.json

@@ -9,7 +9,8 @@
     "editor.formatOnSave": true
   },
   "python.analysis.extraPaths": [
-    "${workspaceFolder}/python"
+    "${workspaceFolder}/python",
+    "${workspaceFolder}/python/spatialdata-codec-writer/scripts"
   ],
   "typescript.tsdk": "node_modules/typescript/lib"
 }

AGENTS.md

@@ -26,5 +26,11 @@ joining work on this repository.
 - Treat layers as independent views of spatial elements: it must be valid for
   multiple layer configs to represent the same underlying element with different
   visual properties, filters, or table-driven encodings.
+- Avoid type assertions (`as ...`) in TypeScript when a library overload,
+  local type guard, schema parser, discriminated union, or narrower API
+  contract can express the same fact. If an assertion is unavoidable at an
+  external boundary (for example an untyped WASM module or a TypeScript
+  correlation limitation), keep it local and add a short comment explaining why
+  the compiler cannot prove it.
 - Worktrees share `.git` but not working state. Documents intended to outlive
   the current branch must land on `main`.

docs/docs/vis/codec-fixtures.mdx

@@ -0,0 +1,191 @@
+---
+sidebar_position: 6
+---
+
+# Codec fixture guide
+
+SpatialData codec fixtures live under `test-fixtures/codecs/` and are generated
+separately from the versioned `blobs()` fixtures. The reference fixtures are
+small OME-Zarr/SpatialData images using:
+
+- `imagecodecs_jpeg2k` (registered JP2K codec id)
+- `experimental.openjph_htj2k` (OpenJPH WASM HTJ2K codec id)
+
+Generated stores include:
+
+- `jpeg2k.zarr` — small multiscale Mandelbrot (`[1,1,1,64,64]`)
+- `mandelbulb.zarr` — multi-dimensional Mandelbulb volume (`[2,1,8,128,128]`, HTJ2K, one plane per t/z chunk)
+- `htj2k.zarr` / `htj2k-demo.zarr` — HTJ2K smoke and encode-demo stores (when the WASM encoder is available)
+
+Older fixtures may use `experimental.imagecodecs_htj2k`; the frontend decodes both.
+
+## Try it in the browser
+
+From the repo root:
+
+```bash
+pnpm test:fixtures:generate:codecs
+pnpm --filter @spatialdata/vis dev
+```
+
+Open [http://127.0.0.1:5173/codec](http://127.0.0.1:5173/codec).
+
+The route calls `enableWorkerChunkDecode()` from `zarrextra/workers`, which uses
+[`@fideus-labs/fizarrita`](https://www.npmjs.com/package/@fideus-labs/fizarrita)
+with a custom codec worker that registers JP2K and experimental HTJ2K support
+inside the worker before decode. Use the codec selector on the page to switch between `/test-fixtures/codecs/jpeg2k.zarr`,
+`/test-fixtures/codecs/mandelbulb.zarr`, and `/test-fixtures/codecs/htj2k-demo.zarr`.
+When HTJ2K is selected, the demo loads a single `htj2k-demo.zarr` store containing
+`mandelbrot_lossless`, `mandelbrot_balanced`, and `mandelbrot_small` image layers
+(512×512 Mandelbrot, two pyramid levels). Use the `SpatialCanvas` layer controls to
+toggle between presets. Manifests are at
+`/test-fixtures/codecs/jpeg2k.manifest.json`,
+`/test-fixtures/codecs/mandelbulb.manifest.json`, and
+`/test-fixtures/codecs/htj2k-encode-demo.manifest.json`.
+
+The vis dev script starts the fixture server on port `38473` and proxies
+`/test-fixtures` through Vite on `5173`. Override the fixture server port with
+`SPATIALDATA_FIXTURE_PORT` if needed.
+
+## Python reference writer
+
+The **publishable** package (`spatialdata-codec-writer`) exposes recompression
+only. Codec test fixtures are generated by repo-local scripts under
+`python/spatialdata-codec-writer/scripts/`.
+
+Useful commands:
+
+```bash
+pnpm test:fixtures:generate:codecs
+uv run --directory python/spatialdata-codec-writer spatialdata-codec-writer inspect ../../test-fixtures/codecs/jpeg2k.manifest.json
+```
+
+### Fixture provenance
+
+Dev fixtures stamp root metadata with an experimental marker rather than a fake
+`spatialdata` library version:
+
+```json
+"spatialdata_attrs": {
+  "experimental_codec_writer": "0.1.0",
+  "note": "codec test fixture; not a spatialdata.write() artifact"
+}
+```
+
+The frontend image schema treats `spatialdata_attrs` as optional on rasters and
+does not use `version` for OME format detection. We intend to propose a
+standardized provenance field to the spatialdata community; feedback welcome.
+
+The sidecar manifest records shape, dtype, chunk shape, codec id, package
+versions, and representative decoded/encoded checksums. The fixture scripts
+validate representative chunks after writing so JS tests can compare decoded
+samples against Python.
+
+### Whole-object recompression
+
+For larger or real examples, use the `recompress` command. It copies the source
+SpatialData store, rewrites configured image rasters with JP2K or experimental
+HTJ2K, and writes labels with Blosc/zstd by default.
+
+```bash
+uv run --directory python/spatialdata-codec-writer spatialdata-codec-writer recompress input.sdata.zarr output-jp2k.zarr --image-key morphology_focus --preset balanced --chunks auto --overwrite
+```
+
+Custom HTJ2K quality (instead of a preset name):
+
+```bash
+uv run --directory python/spatialdata-codec-writer spatialdata-codec-writer recompress \
+  input.sdata.zarr output-htj2k.zarr \
+  --image-key morphology_focus \
+  --codec experimental.openjph_htj2k \
+  --quality 0.001 \
+  --chunks auto \
+  --sibling \
+  --overwrite
+```
+
+`--quality` requires `--image-key` and `--codec experimental.openjph_htj2k`.
+Lower values preserve more detail. With `--sibling`, the output image is named
+like `morphology_focus:htj2k_q0.001`.
+
+For repeatable per-raster settings, use JSON:
+
+```json
+{
+  "default_image": { "codec": "imagecodecs_jpeg2k", "preset": "lossless", "chunks": "auto" },
+  "images": {
+    "morphology_focus": { "preset": "balanced" },
+    "morphology_focus_htj2k_tuned": {
+      "codec": "experimental.openjph_htj2k",
+      "quality": 0.001,
+      "chunks": "auto"
+    },
+    "he_image": { "preset": "small" },
+    "fast_preview": {
+      "codec": "experimental.openjph_htj2k",
+      "preset": "lossless",
+      "chunks": "auto"
+    }
+  },
+  "default_labels": { "codec": "blosc", "clevel": 5 }
+}
+```
+
+JP2K and HTJ2K share preset names (`lossless`, `balanced`, `small`). HTJ2K encode
+uses OpenJPH WASM (`experimental.openjph_htj2k`) via
+`HTJ2KEncoder.setQuality(reversible, quality)`. The quality argument is a float
+quantization factor (lower = higher fidelity, larger output) — not JP2K-style
+0–100. Presets map to `balanced: 0.0002`, `small: 0.001` (calibrated roughly on
+Xenium morphology; see `python/spatialdata-codec-writer/docs/htj2k-wasm-encode-design.md`).
+Override with CLI `--quality 0.001` or per-image JSON `"quality": 0.001` (implies
+lossy unless `"reversible": true`). A future codec-demo UI may add interactive `q`
+exploration on sample regions.
+
+`generate_codec_fixtures.py --experimental-htj2k` also writes
+`htj2k-quality-sweep.manifest.json`, encoding the same Mandelbrot plane at several
+qualities so you can confirm compression responds on detail-rich imagery, and
+`htj2k-encode-demo.manifest.json` with three browser-viewable multiscale image
+layers in one store (`htj2k-demo.zarr`: `mandelbrot_lossless`, `mandelbrot_balanced`,
+`mandelbrot_small`) at 512×512 with 64×64 chunks. The main `htj2k.zarr` fixture
+remains a small 64×64 Mandelbrot raster for fast CI smoke tests.
+`mandelbulb.zarr` exercises multi `t`/`z` HTJ2K decode in JS (`images/mandelbulb`, shape
+`[2, 1, 8, 128, 128]`, chunks `(1, 1, 1, 128, 128)`).
+
+Browser-targeted image codec output is limited to `uint8`, `int8`, `uint16`, and
+`int16`. Labels are not image-codec-compressed in v1; they stay lossless through
+Blosc/zstd because integer IDs and wider dtypes such as `uint32` are not
+considered supported by the current JavaScript decoder paths.
+
+Useful Xenium morphology experiment:
+
+```bash
+uv run --directory python/spatialdata-codec-writer spatialdata-codec-writer recompress /path/to/input.zarr /path/to/output.zarr --image-key morphology_focus --preset balanced --chunks auto --overwrite
+```
+
+When serving recompressed stores for browser experiments, disable HTTP caching
+or use a fresh output path after each rewrite. `http-server` defaults to
+cacheable responses, so overwriting a Zarr store in place can leave the browser
+with stale array metadata that points at the old codec while the chunk bytes have
+changed.
+
+```bash
+bunx http-server --cors -c-1 /private/tmp
+```
+
+Repo scripts may set `UV_CACHE_DIR=.tmp/uv-cache` for sandbox or CI cache
+isolation, but normal user-facing `uv run` commands do not need it.
+
+## JS codec flow
+
+`zarrextra` exposes:
+
+- `registerJpeg2kCodec()` for the registered `imagecodecs_jpeg2k` id (Node/CI).
+- `enableWorkerChunkDecode()` from `zarrextra/workers` for browser apps (fizarrita
+  worker pool + custom codec worker with JP2K and experimental HTJ2K registration).
+- `registerExperimentalHtj2kCodec()` for OpenJPH HTJ2K decode (`experimental.openjph_htj2k`
+  and legacy `experimental.imagecodecs_htj2k`) in Node/CI smoke tests.
+- `loadOmeZarrMultiscalesFromStore()` for loading multiscales from a
+  `zarr.Readable` store without going back through Viv's URL loader.
+
+`@spatialdata/vis` now prefers raster element stores (`images/<key>` and
+`labels/<key>`) and keeps Viv URL loading as a compatibility fallback.

docs/docs/vis/headless-viewer.mdx

@@ -212,6 +212,9 @@ Open **http://127.0.0.1:5173/headless** for the headless demo route. It loads
 the local **`blobs.zarr`** fixture (`v0.7.2`) via `SpatialCanvasViewer` with
 externally controlled layer state — no `SpatialCanvas` sidebars.
 
+Open **http://127.0.0.1:5173/codec** for the JP2K codec fixture route. Generate
+it first with `pnpm test:fixtures:generate:codecs`.
+
 The dev script starts the fixture server on port **38473** (override with
 `SPATIALDATA_FIXTURE_PORT`) and proxies `/test-fixtures` through Vite on
 **5173**. The default Sketch UI remains at **http://127.0.0.1:5173/**.

package.json

@@ -20,6 +20,9 @@
     "test:fixtures:generate:0.5.0": "uv run python/scripts/generate_fixtures.py --version 0.5.0",
     "test:fixtures:generate:0.6.1": "uv run python/scripts/generate_fixtures.py --version 0.6.1",
     "test:fixtures:generate:0.7.2": "uv run python/scripts/generate_fixtures.py --version 0.7.2",
+    "test:fixtures:generate:codecs": "node scripts/vendor-openjph-for-python.mjs && UV_CACHE_DIR=.tmp/uv-cache uv run --directory python/spatialdata-codec-writer python scripts/generate_codec_fixtures.py --output-dir ../../test-fixtures/codecs --experimental-htj2k --overwrite",
+    "write-synthetic": "node scripts/vendor-openjph-for-python.mjs && uv run --directory python/spatialdata-codec-writer spatialdata-codec-writer-write-synthetic",
+    "test:codec-writer": "node scripts/vendor-openjph-for-python.mjs && UV_CACHE_DIR=.tmp/uv-cache uv run --directory python/spatialdata-codec-writer pytest",
     "test:server": "node scripts/test-server.js",
     "test:proxy": "node scripts/cors-proxy.js",
     "validate:datasets:js": "node scripts/validate-datasets-js.js",

packages/avivatorish/package.json

@@ -30,6 +30,7 @@
     "@hms-dbmi/viv": "catalog:",
     "@math.gl/core": "catalog:",
     "geotiff": "2.1.4-beta.0",
+    "zarrextra": "workspace:*",
     "zustand": "^5.0.8"
   },
   "devDependencies": {

packages/avivatorish/src/omeZarrMultiscales.ts

@@ -1,7 +1,28 @@
 import { loadOmeZarr } from '@hms-dbmi/viv';
+import { loadOmeZarrMultiscalesFromStore } from 'zarrextra';
+
+export type OmeZarrMultiscalesSource =
+  | string
+  | {
+      url?: string;
+      store?: unknown;
+    };
+
+/** OME-Zarr multiscales pixel sources (no caching). */
+export async function loadOmeZarrMultiscalesData(
+  source: OmeZarrMultiscalesSource
+): Promise<unknown> {
+  if (typeof source !== 'string' && source.store) {
+    return await loadOmeZarrMultiscalesFromStore(
+      source.store as Parameters<typeof loadOmeZarrMultiscalesFromStore>[0]
+    );
+  }
+
+  const url = typeof source === 'string' ? source : source.url;
+  if (!url) {
+    throw new Error('OME-Zarr loading requires either a Zarrita store or a URL.');
+  }
 
-/** OME-Zarr multiscales pixel sources for a URL (no caching). */
-export async function loadOmeZarrMultiscalesData(url: string): Promise<unknown> {
   const loader = await loadOmeZarr(url, { type: 'multiscales' });
   return loader.data;
 }

packages/avivatorish/tsconfig.json

@@ -12,6 +12,11 @@
     "esModuleInterop": true,
     "resolveJsonModule": true,
     "lib": ["ES2020", "DOM", "DOM.Iterable"],
+    "baseUrl": ".",
+    "paths": {
+      "zarrextra": ["../zarrextra/src/index.ts"],
+      "zarrextra/*": ["../zarrextra/src/*"]
+    },
     "types": ["react", "react-dom"]
   },
   "include": ["src", "vite.config.ts", "tests"],

packages/avivatorish/vite.config.ts

@@ -7,7 +7,7 @@ const pkgRoot = fileURLToPath(new URL('.', import.meta.url));
 const baseConfig = defineViteConfig({
   pkgRoot,
   libName: 'SpatialDataAvivatorish',
-  external: ['@hms-dbmi/viv', '@math.gl/core', 'geotiff', /^zustand(?:\/.*)?$/],
+  external: ['@hms-dbmi/viv', '@math.gl/core', 'geotiff', 'zarrita', 'zarrextra', /^zustand(?:\/.*)?$/],
 });
 
 export default mergeConfig(baseConfig, {