Expand description
RS2-Net rodent brain extraction (onnx feature).
RS2-Net is a Swin-UNETR trained within nnU-Net v2 on 1,142 rat and mouse MRIs from 89 centres
(3–17.2 T). It masks rodent brains, where BET’s human-scale surface model and HD-BET (trained on
humans) both fail. This module ports the nnU-Net pipeline RS2-Net ships around the exported
network (rs2-net.onnx, see crate::models), sharing HD-BET’s sliding-window and
resampling helpers (super::hdbet):
- transpose
(z, y, x)→(y, z, x)(the plans’transpose_forward = [1, 0, 2]); - crop to the bounding box of non-zero voxels, z-score normalise over the crop;
- resample linearly to 0.25 × 0.2 × 0.16 mm (nnU-Net’s “separate z” rule for anisotropic voxels: linear in-plane, nearest-neighbour through-plane);
- pad to at least one patch, then Gaussian-weighted sliding-window inference with 50 % overlap (optionally with 8-fold mirroring test-time augmentation);
- resample the logits back the same way, threshold the sigmoid at 0.5 (logit > 0), un-crop and transpose back.
One deliberate deviation: RS2-Net’s export resamples the logits back with the untransposed original spacing (upstream nnU-Net v2 transposes it), so on thick-slice data it picks the wrong axis for nearest-neighbour. This port uses the transposed spacing, as nnU-Net intends; on an in-vivo mouse GRE (0.17 × 0.20 × 0.8 mm) the two masks agree at Dice 0.986. Isotropic data is unaffected.
Swin’s window attention uses a relative position bias only, so the network runs at any patch
that is a multiple of 32 per axis — but the exported graph is traced at a fixed one,
RS2_NET_PATCH: 128 × 96 × 128 instead of RS2-Net’s 128 × 128 × 160, which peaks at
≈4.5 GB and does not fit wasm32’s address space (this one peaks at ≈2.7 GB). On the in-vivo
mouse GRE the two patches agree at Dice 0.987; larger volumes are tiled.
The input should be a magnitude image (the first echo of a multi-echo GRE works well), in the scanner’s native orientation. RS2-Net applies no post-processing.
§Reference pipeline
Rodent (rat, mouse) brains only, at preclinical resolution (the plans resample to 0.25 × 0.2 × 0.16 mm); for human data use BET, HD-BET or MindGrab. No affine is used. The fixed 128 × 96 × 128 patch peaks at ≈2.7 GB, so the network fits a 32-bit WASM heap; larger volumes are tiled with it.
Reference: 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
Reference implementation: https://github.com/VitoLin21/Rodent-Skull-Stripping (GPL-3.0).
Structs§
- Preprocessed 🔒
- nnU-Net-preprocessed volume plus what is needed to map predictions back.
- Rs2Net
Params - RS2-Net inference parameters.
Constants§
- RS2_
NET_ PATCH - Patch the exported graph was traced at, in the transposed
(y, z, x)order. - TARGET_
SPACING 🔒 - nnU-Net target spacing (mm), in the transposed
(y, z, x)order.
Functions§
- postprocess 🔒
- Resample logits back to the cropped grid, threshold, un-crop, transpose back. Returns the
column-major
(nx, ny, nz)mask. - preprocess 🔒
- Transpose, crop to non-zero, z-score, resample.
Noneif the image is entirely zero. - rs2_net
- Rodent brain mask via RS2-Net.
- transpose_
forward 🔒 - Column-major
(nx, ny, nz)— i.e. C-order(z, y, x)— to C-order(y, z, x).