Inference outside Python¶
The Rust binary¶
rust/ builds brainssl-infer, a single executable that runs both models. It
links ONNX Runtime at run time (pass --ort-lib or set ORT_DYLIB_PATH), so the
same build works with a CPU or a GPU build of the runtime.
cargo build --release --manifest-path rust/Cargo.toml
# tumour segmentation of a four-contrast NIfTI (contrasts on the last axis)
brainssl-infer --ort-lib /opt/onnxruntime/lib/libonnxruntime.so segment \
--model exports/segmentation --input scan.nii.gz --output scan_seg.nii.gz \
--probabilities scan_prob.nii.gz
# connectome classification from ROI time series (PCP .1D format)
brainssl-infer --ort-lib /opt/onnxruntime/lib/libonnxruntime.so classify \
--model exports/connectome --timeseries subject_rois_cc200.1D --atlas cc200
Everything the Python pipeline does around the network is reimplemented, not approximated: the midline search, the symmetric crop, per-contrast normalisation, window placement, Gaussian blending, mirror averaging and the region hierarchy for segmentation; the Ledoit-Wolf estimator, the matrix logarithm, the windows and the intercept for classification. All constants are read from the JSON file exported next to each model, so the two implementations cannot drift apart without a test noticing.
Parity¶
scripts/verify_rust_parity.py runs both implementations on the same scan with
the same exported model and compares the midline estimate, the probabilities and
the label maps. On a full 240 x 240 x 155 MSD scan the largest probability
difference was 3.4e-6 and 99.99996% of voxels received the same label (the four
that did not sit exactly on the 0.5 threshold); the Rust binary took 33 s on the
CPU against 79 s for Python with ONNX Runtime. CI repeats the check on every
change, with a small randomly initialised network on synthetic data.
ONNX¶
brainssl export segmentation writes segmentation.onnx for one sliding
window, input image of shape (N, 4, 160, 128, 128) and output logits of the
same spatial size with three channels, and checks the exported graph against
PyTorch before writing. The bilateral blocks export as slices with a negative
step, so no custom operator is needed.
In the browser¶
The dashboard runs the connectome classifier in plain JavaScript inside a Web Worker, on held-out subjects of one cross-validation split, and runs the segmentation network with ONNX Runtime Web (WebGPU when the browser offers it, WebAssembly otherwise), fetching the model from the Hugging Face Hub.
Hugging Face Hub¶
brainssl export all writes two folders with model cards whose numbers are read
from the report. scripts/push_to_hub.py uploads them as
<user>/brainssl-tumour-segmentation and <user>/brainssl-abide-connectome,
and the manual Hugging Face workflow does the same from a GitHub release.