Ballon999/vnet-2d-retinal-layer-onnx

VNet 2D retinal layer — ONNX (browser/WebGPU)

4-class retinal layer semantic segmentation from H&E eye sections. VNet 2D from keras-unet-collection 0.1.13, trained in Keras 2.15. Exported to ONNX for client-side inference with onnxruntime-web + WebGPU, for the JAX Image Tools viewer.

file	size	notes
model.fp16w.onnx	~590 MB	fp16 weights, fp32 compute — use this
model.json	—	input shape, class count, preprocessing

I/O

Single graph, batch 1, static shapes, NHWC throughout. Softmax is inside the graph, so the client only has to argmax.

tensor	shape	meaning
input	(1, 512, 512, 1)	1 (grayscale — mean of RGB)
vnet_output_activation	(1, 512, 512, 4)	per-pixel class probabilities

Input scaling: 1/255

This is the single easiest thing to get wrong. Nothing in the graph records it and the wrong value does not raise — the model just predicts one class over the entire image. It is not even consistent across this family: the ResUNet-a retinal-layer checkpoints want a caffe mean-subtraction, the VNet one wants 0–1. This model wants 1/255; model.json records it under preprocess.scale.

Use

import * as ort from 'onnxruntime-web';

const meta = await (await fetch(`${base}/model.json`)).json();
const session = await ort.InferenceSession.create(`${base}/model.fp16w.onnx`, {
  executionProviders: ['webgpu'],
});
const out = await session.run({ input: tensor });   // NHWC
// argmax over the last axis -> label map

The graph keeps Keras's NHWC layout, which is the layout canvas pixel data is already in, so preprocessing needs no transpose. Softmax is inside the graph. Reference preprocess / argmax / colorize helpers are in js/segmentation-postprocess.js.

There is no WASM fallback. onnxruntime-web's WASM EP dies with std::bad_alloc on this model at every precision. WebGPU is the only way it runs in a browser, so callers need a real unsupported-browser branch.

Validation

The ONNX graph reproduces the original Keras model at cosine 1.000000 with 100% argmax agreement. mIoU 0.9054 against the held-out masks shipped alongside the checkpoint, with predicted class areas tracking ground truth closely ([145253, 56522, 51745, 8624] vs [148594, 55844, 48880, 8826]). On WebGPU the fp16w label map agrees with the fp32 reference to 99.9996% at an identical mIoU, in ~536 ms per 512² tile (M1 Max).

Why fp16w and not a normal fp16/int8 quantization

model.fp16w.onnx stores weights as fp16 and casts them back to fp32, so every kernel still runs in fp32. It halves the download without touching the arithmetic.

That indirection is necessary because onnxruntime-web's WebGPU EP accumulates natively in f16, while onnxruntime's CPU/WASM backends quietly upcast to fp32 — so a genuinely fp16 graph validates cleanly offline and then degrades in a browser, silently. Measured on ResUNet-a 2D, agreement with the fp32 label map:

variant	size	CPU	WebGPU
fp16 weights / fp32 math	1×	99.995%	99.9947%
true fp16	1×	99.990%	87.10%
static int8 QDQ	0.5×	99.296%	26.83%

Unchanged on onnxruntime-web 1.27.0. Always diff a new model against a CPU reference in a browser — the WebGPU EP returns wrong numbers rather than raising.

License / provenance

Weights © The Jackson Laboratory. Architecture from keras-unet-collection 0.1.13 (MIT), trained with Keras 2.15 / Python 3.10.

Reproducing the export needs Python 3.10 — keras-unet-collection builds its ASPP resize steps as Lambda layers, which Keras 2 serialises as marshalled CPython bytecode; that is not portable across Python minor versions, so 3.11+ fails with bad marshal data (unknown type code).

Export tooling: browser-onnx-tools (export/export_keras_onnx.py), which also applies the WebGPU graph fixes described below.