diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 0dd43e9e27..0c6d1ad621 100644 --- a/.github/workflows/build.yml +++ b/.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 diff --git a/src/async_computation/decode_jxl.ts b/src/async_computation/decode_jxl.ts new file mode 100644 index 0000000000..916e9567fa --- /dev/null +++ b/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] }; + }, +); diff --git a/src/async_computation/decode_jxl_request.ts b/src/async_computation/decode_jxl_request.ts new file mode 100644 index 0000000000..9f1c2fe358 --- /dev/null +++ b/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"); diff --git a/src/datasource/precomputed/async_computation.ts b/src/datasource/precomputed/async_computation.ts index afd324330e..e6b6457a4f 100644 --- a/src/datasource/precomputed/async_computation.ts +++ b/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"; diff --git a/src/datasource/precomputed/backend.ts b/src/datasource/precomputed/backend.ts index 3178907329..ac308f0db0 100644 --- a/src/datasource/precomputed/backend.ts +++ b/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( diff --git a/src/datasource/precomputed/base.ts b/src/datasource/precomputed/base.ts index c3414e5e9a..ee925af6db 100644 --- a/src/datasource/precomputed/base.ts +++ b/src/datasource/precomputed/base.ts @@ -27,6 +27,7 @@ export enum VolumeChunkEncoding { COMPRESSED_SEGMENTATION = 2, COMPRESSO = 3, PNG = 4, + JXL = 5, } export class VolumeChunkSourceParameters { diff --git a/src/sliceview/backend_chunk_decoders/jxl.ts b/src/sliceview/backend_chunk_decoders/jxl.ts new file mode 100644 index 0000000000..914367a715 --- /dev/null +++ b/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); +} diff --git a/src/sliceview/jxl/Cargo.toml b/src/sliceview/jxl/Cargo.toml new file mode 100644 index 0000000000..344f67899f --- /dev/null +++ b/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 diff --git a/src/sliceview/jxl/Dockerfile b/src/sliceview/jxl/Dockerfile new file mode 100644 index 0000000000..d7adaf55da --- /dev/null +++ b/src/sliceview/jxl/Dockerfile @@ -0,0 +1,5 @@ +FROM rust:slim-bullseye + +RUN rustup target add wasm32-unknown-unknown + + diff --git a/src/sliceview/jxl/build.sh b/src/sliceview/jxl/build.sh new file mode 100755 index 0000000000..674e31e4fe --- /dev/null +++ b/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 diff --git a/src/sliceview/jxl/build_wasm.sh b/src/sliceview/jxl/build_wasm.sh new file mode 100755 index 0000000000..1bc9daf491 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/sliceview/jxl/index.ts b/src/sliceview/jxl/index.ts new file mode 100644 index 0000000000..b41f3fba84 --- /dev/null +++ b/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 | 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 { + 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); + } + } +} diff --git a/src/sliceview/jxl/jxl_decoder.wasm b/src/sliceview/jxl/jxl_decoder.wasm new file mode 100755 index 0000000000..d5c3b161c4 Binary files /dev/null and b/src/sliceview/jxl/jxl_decoder.wasm differ diff --git a/src/sliceview/jxl/src/lib.rs b/src/sliceview/jxl/src/lib.rs new file mode 100644 index 0000000000..5b9bbf63a6 --- /dev/null +++ b/src/sliceview/jxl/src/lib.rs @@ -0,0 +1,147 @@ +use std::ptr; +use std::alloc::{alloc, dealloc, Layout}; +use std::slice; + +use jxl_oxide::{FrameBuffer, JxlImage, PixelFormat}; + +#[no_mangle] +pub fn malloc(size: usize) -> *mut u8 { + let layout = Layout::from_size_align(size, std::mem::align_of::()).unwrap(); + unsafe { + let ptr = alloc(layout); + if ptr.is_null() { + panic!("Memory allocation failed"); + } + ptr + } +} + +#[no_mangle] +pub fn free(ptr: *mut u8, size: usize) { + let layout = Layout::from_size_align(size, std::mem::align_of::()).unwrap(); + unsafe { + dealloc(ptr, layout); + } +} + +#[no_mangle] +pub fn width(ptr: *mut u8, input_size: usize, output_size: usize) -> i32 { + if ptr.is_null() || input_size == 0 || output_size == 0 { + return -1; + } + + let data: &[u8] = unsafe { + slice::from_raw_parts(ptr, input_size) + }; + + let image = match JxlImage::builder().read(data) { + Ok(image) => image, + Err(_image) => return -2, + }; + + for keyframe_idx in 0..image.num_loaded_keyframes() { + let frame = match image.render_frame(keyframe_idx) { + Ok(frame) => frame, + Err(_frame) => return -3, + }; + + let stream = frame.stream(); + return stream.width() as i32; + } + + -4 as i32 +} + +#[no_mangle] +pub fn height(ptr: *mut u8, input_size: usize, output_size: usize) -> i32 { + if ptr.is_null() || input_size == 0 || output_size == 0 { + return -1; + } + + let data: &[u8] = unsafe { + slice::from_raw_parts(ptr, input_size) + }; + + let image = match JxlImage::builder().read(data) { + Ok(image) => image, + Err(_image) => return -2, + }; + + for keyframe_idx in 0..image.num_loaded_keyframes() { + let frame = match image.render_frame(keyframe_idx) { + Ok(frame) => frame, + Err(_frame) => return -3, + }; + + let stream = frame.stream(); + return stream.height() as i32; + } + + -4 as i32 +} + +#[no_mangle] +pub fn decode(ptr: *mut u8, input_size: usize, output_size: usize) -> *const u8 { + if ptr.is_null() || input_size == 0 || output_size == 0 { + return ptr::null(); + } + + let data: &[u8] = unsafe { + slice::from_raw_parts(ptr, input_size) + }; + + let image = match JxlImage::builder().read(data) { + Ok(image) => image, + Err(_image) => return std::ptr::null_mut(), + }; + + let mut output_buffer = Vec::with_capacity(output_size); + + for keyframe_idx in 0..image.num_loaded_keyframes() { + let frame = match image.render_frame(keyframe_idx) { + Ok(frame) => frame, + Err(_frame) => return std::ptr::null_mut(), + }; + + let mut stream = frame.stream(); + let mut fb = FrameBuffer::new( + stream.width() as usize, + stream.height() as usize, + stream.channels() as usize, + ); + stream.write_to_buffer(fb.buf_mut()); + + match image.pixel_format() { + PixelFormat::Gray => { + for pixel in fb.buf() { + let value = (pixel * 255.0).clamp(0.0, 255.0) as u8; + output_buffer.push(value); + } + }, + PixelFormat::Rgb => { + for pixel in fb.buf() { + let value = (pixel * 255.0).clamp(0.0, 255.0) as u8; + output_buffer.push(value); + } + } + PixelFormat::Rgba => { + for pixel in fb.buf() { + let value = (pixel * 255.0).clamp(0.0, 255.0) as u8; + output_buffer.push(value); + output_buffer.push(255); // Alpha channel set to fully opaque + } + } + _ => return std::ptr::null_mut(), + } + } + + // Allocate memory in WASM and return a pointer and length + let ptr = output_buffer.as_ptr(); + + // Ensure that the memory is not dropped until after we return + std::mem::forget(output_buffer); + + ptr +} + + diff --git a/src/sliceview/png/index.ts b/src/sliceview/png/index.ts index 335937304f..c138aca232 100644 --- a/src/sliceview/png/index.ts +++ b/src/sliceview/png/index.ts @@ -71,13 +71,13 @@ function readHeader(buffer: Uint8Array): { } if (buffer.length < 8 + 4) { - throw new Error("png: Invalid image size: {buffer.length}"); + throw new Error(`png: Invalid image size: ${buffer.length}`); } // check for header for magic sequence const validMagic = arrayEqualTrucated(magicSpec, buffer); if (!validMagic) { - throw new Error(`png: didn't match magic numbers: {buffer.slice(0,8)}`); + throw new Error(`png: didn't match magic numbers: ${buffer.slice(0, 8)}`); } // offset into IHDR chunk so we can read more naturally @@ -86,7 +86,7 @@ function readHeader(buffer: Uint8Array): { const chunkHeaderLength = 12; // len (4), code (4), CRC (4) if (buffer.length < magicSpec.length + chunkLength + chunkHeaderLength) { - throw new Error("png: Invalid image size: {buffer.length}"); + throw new Error(`png: Invalid image size: ${buffer.length}`); } const chunkCode = [4, 5, 6, 7].map((i) =>