ImageGrains 2.0: Browser-Ready ONNX Models (WebGPU & WASM)

Self-contained ONNX exports of the official ImageGrains 2.0 pretrained weights for instance segmentation of sedimentary grains (fluvial gravels, cobbles, pebbles, and clasts) to derive Particle Size Distributions ($D_{16}, D_{50}, D_{84}$) and grain morphometrics.

These exports strip out Python/Numba dependencies and pack all model weights into single .onnx protobuf files ($< 2\text{ GB}$ browser limit) so they can run 100% client-side in the browser via onnxruntime-web (WebGPU / WebAssembly) with zero server uploads.

Available Model Files

File	Upstream Checkpoint	Architecture	Size	Precision	Max Abs Error vs. PyTorch	Target Runtime
`ig2_unet_fp32.onnx`	`IG2_full_set.200525`	Cellpose-2/3 `CPnet` U-Net	$25.3\text{ MB}$	FP32	$3.81 \times 10^{-6}$	WebGPU & CPU WASM (Desktop + Mobile)
`ig2_cpsam_fp16.onnx`	`IG2_full_set_cp_SAM`	Cellpose-SAM ViT-L + Flow Decoder	$\approx 588\text{ MB}$	Native FP16	$3.91 \times 10^{-3}$ ($1\text{ ULP}$)	Desktop WebGPU (`Float16Array` support)

Direct Browser CDN URLs

Fast / Universal U-Net ($25.3\text{ MB}$): https://huggingface.co/inYourOwnBrowser/imagegrains-2.0-onnx/resolve/main/ig2_unet_fp32.onnx
High-Accuracy Cellpose-SAM ($\approx 588\text{ MB}$): https://huggingface.co/inYourOwnBrowser/imagegrains-2.0-onnx/resolve/main/ig2_cpsam_fp16.onnx

Input / Output Tensor Contracts

Both models output the exact same 3-channel spatial prediction tensor ([flow_y, flow_x, cellprob]), allowing them to share identical downstream Euler flow-tracking and mask-extraction post-processing.

1. `ig2_unet_fp32.onnx` (Fast / Universal)

Input Name: image
Input Shape & Dtype: [batch, 2, height, width] (float32 / Float32Array) — supports dynamic spatial dimensions (e.g., $256 \times 256$, $384 \times 512$).
Channel Layout:
- Channel 0: Grayscale image intensity normalized to $[0.0, 1.0]$ (1st–99th percentile scaling recommended).
- Channel 1: All zeros ($0.0$).
Output Name: flows_cellprob
Output Shape & Dtype: [batch, 3, height, width] (float32) — [flow_y, flow_x, cellprob_logits].

2. `ig2_cpsam_fp16.onnx` (High-Accuracy Desktop WebGPU)

Input Name: image
Input Shape & Dtype: [batch, 3, 256, 256] (float16 / Float16Array) — fixed $256 \times 256$ tile resolution; dynamic batch axis ($1$ to $32$).
Channel Layout:
- Channels 0, 1, 2: RGB channels normalized independently to $[0.0, 1.0]$ using 1st–99th percentile intensity clipping.
Output Name: flows_cellprob
Output Shape & Dtype: [batch, 3, 256, 256] (float16) — [flow_y, flow_x, cellprob_logits].

Quickstart: Browser Inference (`onnxruntime-web`)

import * as ort from 'onnxruntime-web/webgpu';

// 1. Load the 25.3 MB U-Net model directly from Hugging Face CDN (WebGPU + WASM fallback)
const MODEL_URL =
  'https://huggingface.co/inYourOwnBrowser/imagegrains-2.0-onnx/resolve/main/ig2_unet_fp32.onnx';

const session = await ort.InferenceSession.create(MODEL_URL, {
  executionProviders: ['webgpu', 'wasm'],
});

// 2. Prepare a 256x256 tile: Ch 0 = normalized grayscale [0,1], Ch 1 = zeros
const tileData = new Float32Array(1 * 2 * 256 * 256);
// ... populate tileData[0 .. 256*256 - 1] with normalized grayscale pixels ...

const inputTensor = new ort.Tensor('float32', tileData, [1, 2, 256, 256]);
const results = await session.run({ image: inputTensor });
const flowsCellprob = results.flows_cellprob.data; // Float32Array [3, 256, 256]

For ig2_cpsam_fp16.onnx, fetch from https://huggingface.co/inYourOwnBrowser/imagegrains-2.0-onnx/resolve/main/ig2_cpsam_fp16.onnx, pass executionProviders: ['webgpu'], and construct inputs with new Float16Array(1 * 3 * 256 * 256) (new ort.Tensor('float16', tileData, [1, 3, 256, 256]), supported in Chrome/Edge $\ge 135$ and Safari $\ge 17.4$).

Export Details & Numerical Parity

ig2_unet_fp32.onnx ($25.3\text{ MB}$): Exported from cellpose==3.1.1.1 (ONNX opset 17). Replaced CPnet.make_style's dynamic-kernel avg_pool2d with an equivalent ReduceMean operation (x.mean(dim=(-2, -1))) to enable arbitrary dynamic height and width inputs. Verified against PyTorch FP32 ($\text{max_abs_error} = 3.81 \times 10^{-6}$).
ig2_cpsam_fp16.onnx ($\approx 588\text{ MB}$): Exported from cellpose==4.1.x via torch.export.export + PyTorch Dynamo ONNX exporter (opset 18). Instantiated natively in torch.float16 prior to tracing to prevent mixed-precision graph mismatches, then consolidated into a single protobuf binary. Verified against PyTorch FP16 ($\text{max_abs_error} = 3.91 \times 10^{-3}$, matching $1\text{ ULP}$ in float16).

Upstream Resources & Data Stores

Original GitHub Repository: dmair1989/imagegrains
Official Pretrained Weights (Zenodo): DOI: 10.5281/zenodo.15728186
Cellpose Framework: MouseLand/cellpose

Citation & Attribution

If you use these models in academic work, reports, or software, please cite the original ImageGrains and Cellpose publications:

@article{Mair2024ImageGrains,
  author  = {Mair, David and Witz, Guillaume and Do Prado, Ariel Henrique and Garefalakis, Philippos and Schlunegger, Fritz},
  title   = {Automated detecting, segmenting and measuring of grains in images of fluvial gravels: The potential for large and precise data from specialist deep learning models and transfer learning},
  journal = {Earth Surface Processes and Landforms},
  volume  = {49},
  number  = {3},
  pages   = {1099--1116},
  year    = {2024},
  doi     = {10.1002/esp.5755}
}

@dataset{Mair2025ImageGrains2Weights,
  author    = {Mair, David and Do Prado, Ariel Henrique and Witz, Guillaume and Schlunegger, Fritz},
  title     = {ImageGrains 2.0 Pretrained Models and Weights},
  year      = {2025},
  publisher = {Zenodo},
  doi       = {10.5281/zenodo.15728186},
  url       = {https://zenodo.org/records/15728186}
}

@article{Stringer2021Cellpose,
  author  = {Stringer, Carsen and Wang, Tim and Michaelos, Michalis and Pachitariu, Marius},
  title   = {Cellpose: a generalist algorithm for cellular segmentation},
  journal = {Nature Methods},
  volume  = {18},
  number  = {1},
  pages   = {100--106},
  year    = {2021},
  doi     = {10.1038/s41592-020-01018-x}
}

License

Model Weights: CC BY 4.0 (derived from 10.5281/zenodo.15728186).
Upstream imagegrains Software: GNU General Public License v3.0 (GPL-3.0).
Upstream cellpose Architecture: BSD-3-Clause.

Downloads last month: -; Downloads are not tracked for this model. How to track