Skip to main content

Module rs2net

Module rs2net 

Source
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):

  1. transpose (z, y, x) → (y, z, x) (the plans’ transpose_forward = [1, 0, 2]);
  2. crop to the bounding box of non-zero voxels, z-score normalise over the crop;
  3. 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);
  4. pad to at least one patch, then Gaussian-weighted sliding-window inference with 50 % overlap (optionally with 8-fold mirroring test-time augmentation);
  5. 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.
Rs2NetParams
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. None if 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).