jax-ai-js

Browser-side inference for JAX image-analysis models, on WebGPU. Two ports, each mirroring a Temporal worker minus the Temporal:

| module | worker | in → out | | ---------- | ---------------------- | ---------------------------------------- | | retinal/ | retinal-layer-worker | image → layer outlines (ONL / INL / GCL) | | stain/ | stain-adapter-worker | image → virtually H&E-stained image |

Models are ONNX exports hosted on the Hugging Face Hub, produced by browser-onnx-tools.

Status: both engines are complete and covered by tests, but they have not yet been exercised end-to-end in a browser against the live Hub models — every number quoted below was measured in Python during export, or in unit tests against fixtures generated from it. Treat the first in-browser run as the real integration test.

Browser requirements

• WebGPU ('gpu' in navigator) — Chrome ≥113, Safari ≥18. There is no WASM fallback: onnxruntime-web's WASM EP dies with std::bad_alloc on these graphs at 512², at every precision. Hosts get WebGpuRequiredError and need a real unsupported-browser branch.
• onnxruntime-web ~1.26.0 as a peer dependency, so the host controls its version and asset pipeline.
• Enough VRAM for the model: the retinal VNet is 590 MB, the stain generators 109 MB each.

npm install jax-ai-js onnxruntime-web

Retinal layer segmentation

import { RetinalSegmenter } from 'jax-ai-js';

const seg = await RetinalSegmenter.fromPretrained({
  modelUrl: 'https://huggingface.co/Ballon999/vnet-2d-retinal-layer-onnx/resolve/main/model.fp16w.onnx',
  fetcher: myAuthenticatedFetch,       // optional; see "Fetching" below
  onProgress: (loaded, total) => …,
});

const { regions, classAreas, unassignedFraction } = await seg.segment(rgbaImage, {
  onProgress: (f) => …,
  signal: abortController.signal,
});

regions are traced polygons (exterior + holes, image pixel coordinates) with classId / className. Background (class 0) is excluded by default — pass classFilter to change that.

Stain adaptation

import { StainAdapter, STAIN_MODELS } from 'jax-ai-js';

const adapter = await StainAdapter.fromPretrained({ modelUrl: STAIN_MODELS[0].modelUrl });
const { image, clampedFraction } = await adapter.restain(rgbaImage);

Output is RGBA at the same dimensions as the input. Match the checkpoint's magnification (20x / 40x) to the level you read pixels at, or downsample first — a 20x model on 40x pixels sees everything at twice the expected scale.

Stain normalization

Moved to stain-normalization-js, which also gained the two staintools methods this package never had:

| method | | | ------------------------- | -------------------------------------------------- | | l_histogram, reinhard | colour statistics — moved from here unchanged | | macenko, vahadane | new — stain separation, ported from staintools |

Re-exported from here for the two that moved, so existing imports keep working. Import the package directly for stain separation, the WebGPU path, and precomputed fitted state:

import { createNormalizer, parseFittedState } from 'stain-normalization-js';

const norm = createNormalizer('macenko');
norm.fit(referenceImage);
const evened = norm.transform(tile);

The JAX normalization targets are published as precomputed fitted state at Ballon999/stain-adapter-normalization-targets — fitting needs the reference at full resolution, and one of them is 170 megapixels.

Things worth knowing before you trust the output

WebGPU is required, and there is no fallback. onnxruntime-web's WASM EP dies with std::bad_alloc on these graphs at 512², at every precision — so this is not a "slow path" situation. Hosts need a real unsupported-browser branch; WebGpuRequiredError is thrown to make it explicit.

Read model.json; do not infer. Both families carry conventions that are invisible in the ONNX graph and silent when wrong:

• Input scale (retinal): VNet wants x/255, the ResUNet-a checkpoints want raw 0–255. The wrong one makes the model predict a single class over the whole image — mIoU collapses from 0.906 to 0.177, with no error. parseRetinalMeta refuses to guess.
• Channel packing (stain): gr_2empty ([grey,0,0]) and triple_gray ([grey,grey,grey]) are the same shape, so a swap is undetectable downstream. It costs up to 17 dB. parseStainMeta cross-checks the packing against the input channel count.

Colour conversion is matched to skimage, not to the textbook. The stain checkpoints were trained through skimage.color, so color.ts reproduces skimage specifically — its rounded 1996 sRGB matrix rather than the precise primaries, and its rounded Lab constants (0.008856 / 7.787) rather than the exact CIE rationals. Both "corrections" are more accurate in the abstract and wrong here. Verified to 1.4e-14 against skimage 0.22 over 700 random colours.

Greyscale is BT.709 (skimage.color.rgb2gray), not BT.601. Measured on the 20x checkpoint against ground truth: BT.709 33.30 dB, BT.601 28.59 dB, channel mean 22.49 dB — all three produce a perfectly plausible image.

Clamping, not wrapping. The reference Python casts generator output to uint8 without clipping, so overshoot past [-1,1] wraps modulo 256 and turns bright pixels black. This clamps instead, and reports clampedFraction so a divergence from a server-side run is visible rather than mysterious.

Two deliberate divergences from the workers

Both are covered by tests, so neither can be "fixed" back by accident.

Edge coverage. retinal-layer-worker computes (width // 512) * 512 and crops the remainder, so on a 1864×1438 image a 328px strip down the right and 414px along the bottom are never predicted. Here the last row/column shift flush to the edge instead, and every pixel is covered.

Greyscale for the retinal VNet. The worker calls cvtColor(img, COLOR_RGB2GRAY) on data OpenCV decoded as BGR, so red and blue coefficients are swapped. Almost certainly unintentional — but it is what the deployed server does, so it is available as grayscale: 'worker-bgr'. The default is 'mean', which is what the export's mIoU 0.906 was measured with.

Fetching

Every network call goes through an injectable fetcher. Pass the host's own HTTP stack when models sit behind auth — a raw fetch bypasses an Angular auth interceptor, 302s to a login page, and then fails CORS. Downloads are cached in IndexedDB, keyed by URL plus ETag.

Development

npm install
npm test          # vitest
npm run typecheck
npm run lint
npm run build

License

MIT © The Jackson Laboratory. Model weights are licensed separately — see each Hub repo's model card.