sbollmann/rs2-net-onnx

RS2-Net (ONNX export) — rodent brain extraction

An ONNX export of RS2-Net, a Swin-UNETR trained within nnU-Net v2 on 1,142 rat and mouse brain MRIs from 89 centres. It maps a magnitude image to a brain mask, and is used by QSM.rs (qsm_core::bet::rs2_net) and QSMbly. It is staged here for astewartau/QSM.rs#126 to be mirrored into qsmxt/qsm-onnx-weights.

File	Bytes	SHA-256
rs2-net.onnx	63,456,162	a120ce43b06a9f3ddecf85bf18f281cbb54c354ef15b34a6bcaaffa4e3c4a570

Provenance

Exported from the authors' released checkpoint (RS2_pretrained_model.pt, release main-tag of VitoLin21/Rodent-Skull-Stripping) with scripts/onnx-export/export_rs2net.py from QSM.rs#126 (Python 3.10, torch 2.3.0 CPU, MONAI 1.3.0, onnx 1.23, opset 17). The export is reproducible: re-running it gives identical bytes. The network is unchanged; two export fixes are described in the script: the _orig_mod. state-dict prefix is stripped, and a traced layer_norm shape is made static.

• Input input: float32 [1, 1, 128, 96, 128], a fixed patch in nnU-Net's transposed (y, z, x) order. The Swin window padding is baked in at export time, so this is the only size it accepts. RS2-Net's own 128×128×160 patch needs ~4.5 GB, more than 32-bit WebAssembly can address; this one peaks at ~2.7 GB.
• Output logits: float32 [1, 1, 128, 96, 128]; the brain is where sigmoid(logit) > 0.5.
• The network expects RS2-Net's nnU-Net preprocessing: transpose, crop to non-zero, z-score, then resample to 0.25×0.2×0.16 mm. Run it through qsm_core::bet::rs2_net, or RS2-Net's own pipeline, rather than on raw volumes.

Verification: against PyTorch the logits agree to max |Δ| 2.2e-3, with the same sign at 99.9998 % of voxels. On an in-vivo mouse GRE, masks from QSM.rs (tract) and from RS2-Net's pipeline running this graph are identical voxel for voxel.

Licence

RS2-Net's code and weights are released under GPL-3.0, and this derived file carries the same licence.

Citation

Lin, Y., Ding, Y., Chang, S., Ge, X., Sui, X., Jiang, Y. (2024). "RS2-Net: An end-to-end deep learning framework for rodent skull stripping in multi-center brain MRI." NeuroImage, 298:120769. https://doi.org/10.1016/j.neuroimage.2024.120769