feat: JPEG-XL Support (encoding: "jxl") (#653)

.github/workflows/build.yml

@@ -296,5 +296,6 @@ jobs:
       - run: ./src/mesh/draco/build.sh
       - run: ./src/sliceview/compresso/build.sh
       - run: ./src/sliceview/png/build.sh
+      - run: ./src/sliceview/jxl/build.sh
       # Check that there are no differences.
       - run: git diff --exit-code

src/async_computation/decode_jxl.ts

@@ -0,0 +1,37 @@
+/**
+ * @license
+ * Copyright 2024 William Silversmith
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { decodeJxl } from "#src/async_computation/decode_jxl_request.js";
+import { registerAsyncComputation } from "#src/async_computation/handler.js";
+import { decompressJxl } from "#src/sliceview/jxl/index.js";
+
+registerAsyncComputation(
+  decodeJxl,
+  async (
+    data: Uint8Array,
+    area: number | undefined,
+    numComponents: number | undefined,
+    bytesPerPixel: number,
+  ) => {
+    const result = await decompressJxl(
+      data,
+      area,
+      numComponents,
+      bytesPerPixel,
+    );
+    return { value: result, transfer: [result.uint8Array.buffer] };
+  },
+);

src/async_computation/decode_jxl_request.ts

@@ -0,0 +1,27 @@
+/**
+ * @license
+ * Copyright 2024 William Silversmith
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import type { DecodedImage } from "#src/async_computation/decode_png_request.js";
+import { asyncComputation } from "#src/async_computation/index.js";
+
+export const decodeJxl =
+  asyncComputation<
+    (
+      data: Uint8Array,
+      area: number | undefined,
+      numComponents: number | undefined,
+      bytesPerPixel: number,
+    ) => DecodedImage
+  >("decodeJxl");

src/datasource/precomputed/async_computation.ts

@@ -1,4 +1,5 @@
 import "#src/async_computation/decode_jpeg.js";
+import "#src/async_computation/decode_jxl.js";
 import "#src/async_computation/decode_gzip.js";
 import "#src/async_computation/decode_compresso.js";
 import "#src/async_computation/decode_png.js";

src/datasource/precomputed/backend.ts

@@ -74,6 +74,7 @@ import { decodeCompressedSegmentationChunk } from "#src/sliceview/backend_chunk_
 import { decodeCompressoChunk } from "#src/sliceview/backend_chunk_decoders/compresso.js";
 import type { ChunkDecoder } from "#src/sliceview/backend_chunk_decoders/index.js";
 import { decodeJpegChunk } from "#src/sliceview/backend_chunk_decoders/jpeg.js";
+import { decodeJxlChunk } from "#src/sliceview/backend_chunk_decoders/jxl.js";
 import { decodePngChunk } from "#src/sliceview/backend_chunk_decoders/png.js";
 import { decodeRawChunk } from "#src/sliceview/backend_chunk_decoders/raw.js";
 import type { VolumeChunk } from "#src/sliceview/volume/backend.js";
@@ -380,6 +381,7 @@ chunkDecoders.set(
 );
 chunkDecoders.set(VolumeChunkEncoding.COMPRESSO, decodeCompressoChunk);
 chunkDecoders.set(VolumeChunkEncoding.PNG, decodePngChunk);
+chunkDecoders.set(VolumeChunkEncoding.JXL, decodeJxlChunk);
 
 @registerSharedObject()
 export class PrecomputedVolumeChunkSource extends WithParameters(

src/datasource/precomputed/base.ts

@@ -27,6 +27,7 @@ export enum VolumeChunkEncoding {
   COMPRESSED_SEGMENTATION = 2,
   COMPRESSO = 3,
   PNG = 4,
+  JXL = 5,
 }
 
 export class VolumeChunkSourceParameters {

src/sliceview/backend_chunk_decoders/jxl.ts

@@ -0,0 +1,39 @@
+/**
+ * @license
+ * Copyright 2016 Google Inc.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { decodeJxl } from "#src/async_computation/decode_jxl_request.js";
+import { requestAsyncComputation } from "#src/async_computation/request.js";
+import { postProcessRawData } from "#src/sliceview/backend_chunk_decoders/postprocess.js";
+import type { VolumeChunk } from "#src/sliceview/volume/backend.js";
+import type { CancellationToken } from "#src/util/cancellation.js";
+
+export async function decodeJxlChunk(
+  chunk: VolumeChunk,
+  cancellationToken: CancellationToken,
+  response: ArrayBuffer,
+) {
+  const chunkDataSize = chunk.chunkDataSize!;
+  const { uint8Array: decoded } = await requestAsyncComputation(
+    decodeJxl,
+    cancellationToken,
+    [response],
+    new Uint8Array(response),
+    chunkDataSize[0] * chunkDataSize[1] * chunkDataSize[2],
+    chunkDataSize[3] || 1,
+    1, // bytesPerPixel
+  );
+  await postProcessRawData(chunk, cancellationToken, decoded);
+}

src/sliceview/jxl/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "jxl-wasm"
+version = "1.0.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[lib]
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+jxl-oxide = "0.9.1"
+
+[profile.release]
+lto = true
+
+[package.metadata.wasm-opt]
+memory = { initial = 10, maximum = 100 }  # Set initial and max memory in MiB

src/sliceview/jxl/Dockerfile

@@ -0,0 +1,5 @@
+FROM rust:slim-bullseye
+
+RUN rustup target add wasm32-unknown-unknown
+
+

src/sliceview/jxl/build.sh

@@ -0,0 +1,13 @@
+#!/bin/bash -xve
+
+# This script builds `jxl_decoder.wasm` using emsdk in a docker container.
+
+cd "$(dirname "$0")"
+
+docker build .
+docker run \
+       --rm \
+       -v ${PWD}:/src \
+       -u $(id -u):$(id -g) \
+       $(docker build -q .) \
+       /src/build_wasm.sh

src/sliceview/jxl/build_wasm.sh

@@ -0,0 +1,6 @@
+#!/bin/bash -xve
+
+cd /src
+cargo build --target wasm32-unknown-unknown --release
+cp /src/target/wasm32-unknown-unknown/release/jxl_wasm.wasm /src/jxl_decoder.wasm
+rm -r target

src/sliceview/jxl/index.ts

@@ -0,0 +1,169 @@
+/**
+ * @license
+ * Copyright 2024 William Silvermsith
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import type { DecodedImage } from "#src/async_computation/decode_png_request.js";
+
+const libraryEnv = {};
+
+let jxlModulePromise: Promise<WebAssembly.Instance> | undefined;
+
+async function getJxlModulePromise() {
+  if (jxlModulePromise === undefined) {
+    jxlModulePromise = (async () => {
+      const m = (
+        await WebAssembly.instantiateStreaming(
+          fetch(new URL("./jxl_decoder.wasm", import.meta.url)),
+          {
+            env: libraryEnv,
+            wasi_snapshot_preview1: libraryEnv,
+          },
+        )
+      ).instance;
+      return m;
+    })();
+  }
+  return jxlModulePromise;
+}
+
+// header constants
+// obtained from
+// // https://github.com/libjxl/libjxl/blob/8f22cb1fb98ed27ceee59887bd291ef4d277c89d/lib/jxl/decode.cc#L118-L130
+const magicSpec = [
+  0,
+  0,
+  0,
+  0xc,
+  "J".charCodeAt(0),
+  "X".charCodeAt(0),
+  "L".charCodeAt(0),
+  " ".charCodeAt(0),
+  0xd,
+  0xa,
+  0x87,
+  0xa,
+];
+
+// not a full implementation of read header, just the parts we need
+function checkHeader(buffer: Uint8Array) {
+  function arrayEqualTrucated(a: any, b: any): boolean {
+    return a.every((val: number, idx: number) => val === b[idx]);
+  }
+
+  const len = buffer.length;
+  const kCodestreamMarker = 0x0a;
+
+  if (len < 8 + 4) {
+    throw new Error(`jxl: Invalid image size: ${len}`);
+  }
+
+  // JPEG XL codestream: 0xff 0x0a
+  if (len >= 1 && buffer[0] === 0xff) {
+    if (len < 2) {
+      throw new Error(`jxl: Not enough bytes. Got: ${len}`);
+    } else if (buffer[1] === kCodestreamMarker) {
+      // valid codestream
+      return;
+    } else {
+      throw new Error(`jxl: Invalid codestream.`);
+    }
+  }
+
+  // JPEG XL container
+  // check for header for magic sequence
+  const validMagic = arrayEqualTrucated(magicSpec, buffer);
+  if (!validMagic) {
+    throw new Error(`jxl: didn't match magic numbers: ${buffer.slice(0, 12)}`);
+  }
+}
+
+export async function decompressJxl(
+  buffer: Uint8Array,
+  area: number | undefined,
+  numComponents: number | undefined,
+  bytesPerPixel: number,
+): Promise<DecodedImage> {
+  const m = await getJxlModulePromise();
+  checkHeader(buffer);
+
+  area ||= 0;
+  numComponents ||= 1;
+
+  const nbytes = area * bytesPerPixel * numComponents;
+
+  const jxlImagePtr = (m.exports.malloc as Function)(buffer.byteLength);
+  const heap = new Uint8Array((m.exports.memory as WebAssembly.Memory).buffer);
+  heap.set(buffer, jxlImagePtr);
+
+  let imagePtr = null;
+
+  try {
+    const width = (m.exports.width as Function)(
+      jxlImagePtr,
+      buffer.byteLength,
+      nbytes,
+    );
+    const height = (m.exports.height as Function)(
+      jxlImagePtr,
+      buffer.byteLength,
+      nbytes,
+    );
+
+    if (width <= 0 || height <= 0) {
+      throw new Error(
+        `jxl: Decoding failed. Width (${width}) and/or height (${height}) invalid.`,
+      );
+    }
+
+    if (area !== undefined && width * height !== area) {
+      throw new Error(
+        `jxl: Expected width and height (${width} x ${height}, ${width * height}) to match area: ${area}.`,
+      );
+    }
+
+    imagePtr = (m.exports.decode as Function)(
+      jxlImagePtr,
+      buffer.byteLength,
+      nbytes,
+    );
+
+    if (imagePtr === 0) {
+      throw new Error("jxl: Decoding failed. Null pointer returned.");
+    }
+
+    // Likewise, we reference memory.buffer instead of heap.buffer
+    // because memory growth during decompress could have detached
+    // the buffer.
+    const image = new Uint8Array(
+      (m.exports.memory as WebAssembly.Memory).buffer,
+      imagePtr,
+      nbytes,
+    );
+
+    // copy the array so it can be memory managed by JS
+    // and we can free the emscripten buffer
+    return {
+      width: width || 0,
+      height: height || 0,
+      numComponents: numComponents || 1,
+      uint8Array: image.slice(0),
+    };
+  } finally {
+    (m.exports.free as Function)(jxlImagePtr, buffer.byteLength);
+    if (imagePtr) {
+      (m.exports.free as Function)(imagePtr, nbytes);
+    }
+  }
+}