qsm_core/bet/

hdbet.rs

//! HD-BET deep-learning brain extraction (`onnx` feature).
//!
//! HD-BET v2 is an nnU-Net v2 3D U-Net (`PlainConvUNet`, 1 mm isotropic, 96×192×192 patches)
//! trained on 11,751 multi-sequence clinical MRIs; it segments the brain from a single magnitude
//! image. This module is a faithful port of nnU-Net's inference pipeline around the exported
//! network (`hd-bet.onnx`, see [`crate::models`]):
//!
//! 1. crop to the bounding box of non-zero voxels;
//! 2. z-score normalise over the cropped image;
//! 3. resample to 1 mm with cubic splines (nnU-Net's "separate z" rule for anisotropic voxels:
//!    cubic in-plane, nearest-neighbour through-plane);
//! 4. run the network on the whole resampled volume, zero-padded to a size the network accepts
//!    ([`HdBetInference::FullVolume`], the default, as HD-BET v1 did), or — as nnU-Net and
//!    HD-BET v2 do — pad to at least one patch and run Gaussian-weighted sliding-window
//!    inference with 50 % overlap ([`HdBetInference::SlidingWindow`]); either optionally with
//!    8-fold mirroring test-time augmentation;
//! 5. resample the logits back (linear; nearest through-plane for anisotropic voxels), take the
//!    argmax and un-crop.
//!
//! nnU-Net works on SimpleITK arrays in `(z, y, x)` order, which is exactly the memory layout of
//! the crate's column-major `(nx, ny, nz)` volumes — no transposition is needed.
//!
//! The input should be a magnitude image (for multi-echo GRE, the root-sum-of-squares over
//! echoes works well) in roughly standard radiological orientation (axial slices along z), as
//! HD-BET was trained on MNI-aligned data. HD-BET applies no post-processing; combine with
//! [`crate::utils::morphology::fill_holes`] / erosion refinements as needed.
//!
//! # Reference pipeline
//!
//! From the HD-BET repository (MIC-DKFZ/HD-BET 678e44d, v2):
//!
//! - **Contrast.** Trained on pre- and post-contrast T1-w, T2-w and FLAIR; "other sequences will
//!   most likely work as well" (`readme.md:25-28,79-80`). No GRE was in training, and QSM.rs has
//!   no Dice against a reference brain mask on GRE recorded; its validation is agreement with the
//!   `hd-bet` CLI (Dice ≥ 0.99998, `docs/DL_ONNX_SCOPING.md`).
//! - **Orientation.** "Input images must match the orientation of standard MNI152 template! Use
//!   fslreorient2std upfront" (`readme.md:80-82`). Neither HD-BET nor this port reorients: the
//!   array is used in its stored voxel order, and `hd_bet` takes no affine. A host must pass a
//!   volume already in that orientation.
//! - **Resolution.** Any; nnU-Net resamples to 1 mm (step 3 above), so the network's memory
//!   follows the field of view in mm ([`hd_bet_memory_bytes`]).
//! - **Inference.** The v2 CLI uses nnU-Net's sliding window (`tile_step_size=0.5`, Gaussian)
//!   with mirroring test-time augmentation on unless `--disable_tta`
//!   (`HD_BET/hd_bet_prediction.py:28-35`, `HD_BET/entry_point.py:30,48`); the readme
//!   recommends `--disable_tta` on CPU (`readme.md:100-114`). [`HdBetParams::default`] is
//!   whole-volume without TTA.
//! - **WASM.** On a 32-bit WASM heap (4 GiB) the whole-volume default does not fit a typical head
//!   scan; [`HdBetParams::low_memory`] does (1.6 GB on a 184×256×128 mm volume).
//!   [`hd_bet_memory_bytes`] with [`crate::models::Heap::Wasm32`] gives the peak for a grid and
//!   setting.
//!
//! Reference:
//! Isensee, F., Schell, M., Pflueger, I., et al. (2019). "Automated brain extraction of
//! multisequence MRI using artificial neural networks." Human Brain Mapping, 40(17):4952-4964.
//! https://doi.org/10.1002/hbm.24750
//!
//! Reference implementation: https://github.com/MIC-DKFZ/HD-BET (Apache-2.0; weights CC-BY-NC-4.0).
//!
//! # Agreement with the original
//!
//! `tests/models_onnx/separation.rs` runs it on the 2 mm root-sum-of-squares magnitude of QSM-CI's
//! 128×160×96 case (`tests/data/dl_head_small.bin`; a whole head, which HD-BET needs) against its
//! original's verified QSM-CI entry, HD-BET 2.0.1 (`hd-bet-qsmci`) (QSMxT/QSM-CI#316; the fixture
//! is regenerated in Actions by `tests/data/dl_hdbet_qsmci_small.py`): with
//! [`HdBetInference::SLIDING_WINDOW`] and no mirroring, as the original's `hd-bet --disable_tta`
//! runs, the masks are identical (Dice 1.0, 187,055 voxels; measured 2026-10-30).
//!
//! # Pipeline rules
//!
//! The facts above that a configuration can depart from are published as rules in
//! [`crate::pipeline::rules::RULES`] (each with its strength, basis, reason and source): `hd-bet.wasm-heap`.

// The pre/post-processing is plain Rust; only inference needs `onnx`.
#![cfg_attr(not(feature = "onnx"), allow(dead_code))]

use crate::grid::Grid;
#[cfg(feature = "onnx")]
use crate::models::onnx::{OnnxError, OnnxModel, Tensor};
use crate::utils::resample::{resize, resize_axis};

/// How the network is applied to the preprocessed (cropped, 1 mm) volume.
#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum HdBetInference {
    /// One network pass over the whole volume, zero-padded (centred) to the next multiple of
    /// `(32, 32, 16)`. Fastest, and every voxel sees the whole head; peak memory grows with the
    /// field of view in mm (natively 5.7 GB for 184×256×128 mm, 9.3 GB for 240×240×160 mm; see
    /// [`hd_bet_memory_bytes`]).
    FullVolume,
    /// Gaussian-weighted sliding window, as nnU-Net (and the HD-BET v2 CLI) runs it.
    SlidingWindow {
        /// Patch size `(px, py, pz)` in voxels at 1 mm; HD-BET v2 was trained with
        /// `(192, 192, 96)`. Must be multiples of `(32, 32, 16)`. Peak memory follows the patch
        /// (natively 3.4 GB at `(192, 192, 96)`, 1.3 GB at `(128, 128, 64)`; see
        /// [`hd_bet_memory_bytes`]); below ~`(128, 128, 64)` tiles that lie wholly inside the
        /// brain start to be labelled background.
        patch: (usize, usize, usize),
        /// Step between windows as a fraction of the patch (nnU-Net: 0.5).
        tile_step: f64,
    },
}

impl HdBetInference {
    /// nnU-Net's sliding window: HD-BET v2's training patch `(192, 192, 96)`, 50 % overlap.
    pub const SLIDING_WINDOW: Self = Self::SlidingWindow { patch: (192, 192, 96), tile_step: 0.5 };
    /// Memory-bounded sliding window: `(128, 128, 64)` patches, 50 % overlap.
    pub const LOW_MEMORY: Self = Self::SlidingWindow { patch: (128, 128, 64), tile_step: 0.5 };
}

/// HD-BET inference parameters.
#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
#[derive(Clone, Debug, PartialEq)]
pub struct HdBetParams {
    /// Whole-volume (default) or sliding-window inference.
    pub inference: HdBetInference,
    /// 8-fold mirroring test-time augmentation (default off, as in QSM-CI; 8× slower for a
    /// marginal gain).
    pub mirror_tta: bool,
}

impl Default for HdBetParams {
    fn default() -> Self {
        Self { inference: HdBetInference::FullVolume, mirror_tta: false }
    }
}

impl HdBetParams {
    /// Memory-bounded setting for constrained hosts (e.g. 32-bit WASM): `(128, 128, 64)`
    /// sliding-window patches, peak 1.6 GB in a WASM heap on a 184×256×128 mm volume. Agreed
    /// with the native patch at Dice ≈0.99 in validation (and with nnU-Net run at the same patch
    /// size to 2 voxels).
    pub fn low_memory() -> Self {
        Self { inference: HdBetInference::LOW_MEMORY, ..Self::default() }
    }
}

/// nnU-Net target spacing for HD-BET (mm).
const TARGET_SPACING: [f64; 3] = [1.0, 1.0, 1.0];
/// Required divisibility of the patch in nnU-Net `(z, y, x)` order (5 down-samplings, the last
/// one in-plane only).
const PATCH_DIVISOR: [usize; 3] = [16, 32, 32];

/// Brain mask via HD-BET.
///
/// * `magnitude` — magnitude image, column-major `(nx, ny, nz)`.
/// * `model_onnx` — bytes of the exported `hd-bet.onnx` (see [`crate::models`]).
/// * `progress(done, total)` — called after each network evaluation.
///
/// Returns a binary mask (0/1) on the input grid.
#[cfg(feature = "onnx")]
pub fn hd_bet(
    magnitude: &[f64],
    grid: &Grid,
    model_onnx: &[u8],
    params: &HdBetParams,
    progress: impl FnMut(usize, usize),
) -> Result<Vec<u8>, OnnxError> {
    let (nx, ny, nz) = grid.dims;
    assert_eq!(magnitude.len(), nx * ny * nz, "magnitude length must match grid");
    if let HdBetInference::SlidingWindow { patch, tile_step } = params.inference {
        let p = [patch.2, patch.1, patch.0];
        if p.iter().zip(PATCH_DIVISOR).any(|(&p, d)| p == 0 || p % d != 0) {
            return Err(OnnxError::Shape(format!(
                "HD-BET patch {patch:?} must be a non-zero multiple of (32, 32, 16)"
            )));
        }
        if !(tile_step > 0.0 && tile_step <= 1.0) {
            return Err(OnnxError::Shape(format!("tile_step {tile_step} must be in (0, 1]")));
        }
    }
    let Some(pre) = preprocess(magnitude, grid) else {
        return Ok(vec![0; magnitude.len()]);
    };
    let (patch, tile_step) = match params.inference {
        HdBetInference::FullVolume => (full_volume_patch(pre.dims), 1.0),
        HdBetInference::SlidingWindow { patch, tile_step } => ([patch.2, patch.1, patch.0], tile_step),
    };
    let model = OnnxModel::load(model_onnx)?;
    let logits = predict_logits(&pre.data, pre.dims, &model, patch, 2, tile_step, params.mirror_tta, progress)?;
    Ok(postprocess(&logits, &pre))
}

/// nnU-Net-preprocessed volume plus what is needed to map predictions back.
struct Preprocessed {
    /// Cropped, normalised, resampled volume, C-order `(z, y, x)`.
    data: Vec<f32>,
    dims: [usize; 3],
    /// Crop box `[start, end)` per axis in the original `(z, y, x)` volume.
    bbox: [(usize, usize); 3],
    /// Original volume shape `(z, y, x)`.
    full_dims: [usize; 3],
    /// Original spacing `(z, y, x)`.
    spacing: [f64; 3],
}

/// Crop to non-zero, z-score, resample to 1 mm. `None` if the image is entirely zero.
fn preprocess(magnitude: &[f64], grid: &Grid) -> Option<Preprocessed> {
    let (nx, ny, nz) = grid.dims;
    let (vx, vy, vz) = grid.voxel_size;
    let full_dims = [nz, ny, nx];
    let spacing = [vz, vy, vx];

    let bbox = nonzero_bbox(magnitude, full_dims)?;
    let cdims = [bbox[0].1 - bbox[0].0, bbox[1].1 - bbox[1].0, bbox[2].1 - bbox[2].0];
    let mut img = Vec::with_capacity(cdims.iter().product());
    for z in bbox[0].0..bbox[0].1 {
        for y in bbox[1].0..bbox[1].1 {
            let row = (z * ny + y) * nx;
            img.extend_from_slice(&magnitude[row + bbox[2].0..row + bbox[2].1]);
        }
    }

    // ZScoreNormalization without a mask (HD-BET's plans: use_mask_for_norm = false).
    let n = img.len() as f64;
    let mean = img.iter().sum::<f64>() / n;
    let std = (img.iter().map(|v| (v - mean) * (v - mean)).sum::<f64>() / n).sqrt();
    let s = std.max(1e-8);
    img.iter_mut().for_each(|v| *v = (*v - mean) / s);

    let rdims = resampled_dims(cdims, spacing);
    let data = resample_nnunet(&img, cdims, rdims, spacing, TARGET_SPACING, 3)
        .into_iter()
        .map(|v| v as f32)
        .collect();
    Some(Preprocessed { data, dims: rdims, bbox, full_dims, spacing })
}

/// nnU-Net `compute_new_shape`: the `(z, y, x)` dims of a `dims` volume with voxel `spacing` at
/// 1 mm, `round(spacing / target · dims)` with Python's rounding (half to even).
fn resampled_dims(dims: [usize; 3], spacing: [f64; 3]) -> [usize; 3] {
    std::array::from_fn(|a| (spacing[a] / TARGET_SPACING[a] * dims[a] as f64).round_ties_even() as usize)
}

/// The spatial shape `(z, y, x)` of the tensor HD-BET's network sees on a scan on `grid`: with
/// [`HdBetInference::FullVolume`], the field of view resampled to 1 mm and padded to multiples of
/// 16×32×32; with a sliding window, its patch.
///
/// HD-BET crops to the bounding box of the non-zero voxels before resampling. This takes that box
/// to be the whole field of view, as it is for any scan with noise in the background, so on an
/// image whose background was zeroed it is an upper bound.
pub fn hd_bet_network_input(grid: &Grid, params: &HdBetParams) -> [usize; 3] {
    match params.inference {
        HdBetInference::FullVolume => {
            let (nx, ny, nz) = grid.dims;
            let (vx, vy, vz) = grid.voxel_size;
            full_volume_patch(resampled_dims([nz, ny, nx], [vz, vy, vx]))
        }
        HdBetInference::SlidingWindow { patch, .. } => [patch.2, patch.1, patch.0],
    }
}

/// Peak heap of [`hd_bet`] on a scan on `grid`, in bytes, natively or in a `wasm32` heap.
///
/// The network at [`hd_bet_network_input`] ([`crate::models::ModelSpec::peak_bytes`]) plus what
/// the step allocates around it: the 1 mm volume, the padded volume with its two-channel logit
/// and weight accumulators, and the window's buffers (Gaussian, input, prediction, the copy handed
/// to the network). The cropping and resampling before and after the network are counted too,
/// but never set the peak at the resolutions HD-BET is used at. The magnitude passed in and the
/// weight file's bytes are the caller's and are not included. Like [`hd_bet_network_input`], it
/// takes the non-zero bounding box to be the whole field of view.
///
/// Within 0.5–1.5 % above the measured native peak of the whole step, whole-volume and sliding
/// window, at 1 mm, 1.5 mm, 2 mm, 0.6 mm and 0.8×0.8×3 mm (see the [`crate::models`] docs,
/// "Memory").
pub fn hd_bet_memory_bytes(grid: &Grid, params: &HdBetParams, heap: crate::models::Heap) -> u64 {
    let spec = crate::models::find_model("hd-bet").expect("hd-bet is in the registry");
    let (nx, ny, nz) = grid.dims;
    let (vx, vy, vz) = grid.voxel_size;
    let rdims = resampled_dims([nz, ny, nx], [vz, vy, vx]);
    let patch = hd_bet_network_input(grid, params);
    let prod = |d: [usize; 3]| d.iter().map(|&v| v as u64).product::<u64>();
    // The scan, the 1 mm volume, the padded volume, the patch.
    let c = prod([nz, ny, nx]);
    let r = prod(rdims);
    let np = prod(std::array::from_fn(|a| rdims[a].max(patch[a])));
    let pp = prod(patch);
    let big = c.max(r);
    // Cropping and z-scoring in f64, then the spline resize (its copy, and two axis passes alive).
    let pre = 16 * c + 16 * big + 4 * r;
    // The network on the 1 mm volume (f32): the volume; the padded volume, logits (2 channels)
    // and weights on the padded grid; the Gaussian, input, prediction (2) and the network's copy
    // of the input on the patch, and with mirroring the flipped prediction.
    let flips = if params.mirror_tta { 4 * pp } else { 0 };
    let net = spec.peak_bytes(patch, heap) + 4 * r + 16 * np + 20 * pp + flips;
    // Logits back to the scan: the 1 mm volume and logits, one channel in f64 and its resize,
    // both channels on the scan in f64, and the mask.
    let post = spec.resident_bytes() + 20 * r + 16 * big + 16 * c + c;
    pre.max(net).max(post)
}

/// Bounding box `[start, end)` of the non-zero voxels of a C-order `(z, y, x)` volume.
pub(super) fn nonzero_bbox(data: &[f64], dims: [usize; 3]) -> Option<[(usize, usize); 3]> {
    let mut lo = [usize::MAX; 3];
    let mut hi = [0usize; 3];
    let mut any = false;
    for z in 0..dims[0] {
        for y in 0..dims[1] {
            let row = (z * dims[1] + y) * dims[2];
            for x in 0..dims[2] {
                if data[row + x] != 0.0 {
                    any = true;
                    for (a, c) in [z, y, x].into_iter().enumerate() {
                        lo[a] = lo[a].min(c);
                        hi[a] = hi[a].max(c + 1);
                    }
                }
            }
        }
    }
    any.then(|| std::array::from_fn(|a| (lo[a], hi[a])))
}

/// nnU-Net `resample_data_or_seg_to_shape` for one channel (non-segmentation, `order_z = 0`,
/// `force_separate_z = None`): a plain 3D spline resize, unless the voxels are anisotropic
/// (max/min spacing > 3), in which case each slice is resized in-plane with `order` and the
/// low-resolution axis with nearest-neighbour.
pub(super) fn resample_nnunet(
    data: &[f64],
    dims: [usize; 3],
    new_dims: [usize; 3],
    current_spacing: [f64; 3],
    new_spacing: [f64; 3],
    order: usize,
) -> Vec<f64> {
    if dims == new_dims {
        return data.to_vec();
    }
    let Some(axis) = separate_z_axis(current_spacing, new_spacing) else {
        return resize(data, dims, new_dims, order);
    };
    let inplane: Vec<usize> = (0..3).filter(|&a| a != axis).collect();
    let mut out = data.to_vec();
    let mut cur = dims;
    if inplane.iter().any(|&a| dims[a] != new_dims[a]) {
        // skimage.resize on each 2D slice: both in-plane axes (a same-length axis still goes
        // through the spline for order 3, as in scipy), then clip each slice to its own range.
        for &a in &inplane {
            if order == 3 || cur[a] != new_dims[a] {
                (out, cur) = resize_axis(&out, cur, a, new_dims[a], order);
            }
        }
        if order > 1 {
            let range = slice_ranges(data, dims, axis);
            for_each_slice(&mut out, cur, axis, |k, v| *v = v.clamp(range[k].0, range[k].1));
        }
    }
    if cur[axis] != new_dims[axis] {
        (out, _) = resize_axis(&out, cur, axis, new_dims[axis], 0);
    }
    out
}

/// nnU-Net `determine_do_sep_z_and_axis(force_separate_z=None, ...)`: the low-resolution axis if
/// either spacing is anisotropic by more than 3×, and that axis is unique.
pub(super) fn separate_z_axis(current: [f64; 3], new: [f64; 3]) -> Option<usize> {
    let aniso = |s: [f64; 3]| {
        let (mn, mx) = s.iter().fold((f64::INFINITY, 0.0f64), |(a, b), &v| (a.min(v), b.max(v)));
        mx / mn > 3.0
    };
    let lowres = |s: [f64; 3]| {
        let mx = s.iter().copied().fold(0.0f64, f64::max);
        let axes: Vec<usize> = (0..3).filter(|&a| mx / s[a] == 1.0).collect();
        (axes.len() == 1).then(|| axes[0])
    };
    if aniso(current) {
        lowres(current)
    } else if aniso(new) {
        lowres(new)
    } else {
        None
    }
}

/// Per-slice `(min, max)` along `axis` of a C-order volume.
fn slice_ranges(data: &[f64], dims: [usize; 3], axis: usize) -> Vec<(f64, f64)> {
    let mut r = vec![(f64::INFINITY, f64::NEG_INFINITY); dims[axis]];
    let mut idx = 0;
    for i in 0..dims[0] {
        for j in 0..dims[1] {
            for k in 0..dims[2] {
                let s = [i, j, k][axis];
                let v = data[idx];
                r[s] = (r[s].0.min(v), r[s].1.max(v));
                idx += 1;
            }
        }
    }
    r
}

/// Apply `f(slice_index_along_axis, &mut value)` to every voxel of a C-order volume.
fn for_each_slice(data: &mut [f64], dims: [usize; 3], axis: usize, mut f: impl FnMut(usize, &mut f64)) {
    let mut idx = 0;
    for i in 0..dims[0] {
        for j in 0..dims[1] {
            for k in 0..dims[2] {
                f([i, j, k][axis], &mut data[idx]);
                idx += 1;
            }
        }
    }
}

/// The single "patch" covering a whole `(z, y, x)` volume: each axis rounded up to the network's
/// divisor. [`predict_logits`] centres the volume in it and runs one window.
pub(super) fn full_volume_patch(dims: [usize; 3]) -> [usize; 3] {
    std::array::from_fn(|a| dims[a].div_ceil(PATCH_DIVISOR[a]).max(1) * PATCH_DIVISOR[a])
}

/// nnU-Net `compute_steps_for_sliding_window` for one axis.
pub(crate) fn window_steps(size: usize, tile: usize, step: f64) -> Vec<usize> {
    let n = ((size - tile) as f64 / (tile as f64 * step)).ceil() as usize + 1;
    if n == 1 {
        return vec![0];
    }
    let actual = (size - tile) as f64 / (n - 1) as f64;
    (0..n).map(|i| (actual * i as f64).round_ties_even() as usize).collect()
}

/// nnU-Net `compute_gaussian(tile, sigma_scale=1/8, value_scaling_factor=10)`: a Gaussian
/// centred on the patch (σ = tile/8 per axis), peak 10.
pub(crate) fn gaussian_importance(tile: [usize; 3]) -> Vec<f32> {
    let w: [Vec<f64>; 3] = std::array::from_fn(|a| {
        let (c, s) = ((tile[a] / 2) as f64, tile[a] as f64 / 8.0);
        (0..tile[a]).map(|i| (-0.5 * ((i as f64 - c) / s).powi(2)).exp()).collect()
    });
    let mut g = Vec::with_capacity(tile.iter().product());
    for &a in &w[0] {
        for &b in &w[1] {
            for &c in &w[2] {
                g.push((10.0 * a * b * c) as f32);
            }
        }
    }
    g
}

/// Gaussian-weighted sliding-window logits of an nnU-Net-style network with `channels` outputs,
/// channel-major C-order (shape `[channels, dims...]`). `patch` must be a size the network
/// accepts; `tile_step` is the stride as a fraction of the patch; `mirror_tta` averages over the
/// 8 axis mirrorings. Shared with [`super::rs2net`].
#[cfg(feature = "onnx")]
#[allow(clippy::too_many_arguments)]
pub(super) fn predict_logits(
    data: &[f32],
    dims: [usize; 3],
    model: &OnnxModel,
    patch: [usize; 3],
    channels: usize,
    tile_step: f64,
    mirror_tta: bool,
    mut progress: impl FnMut(usize, usize),
) -> Result<Vec<f32>, OnnxError> {
    // Pad (centred, extra voxel on the high side) to at least one patch, with zeros (= the mean
    // after normalisation), as pad_nd_image does.
    let pdims: [usize; 3] = std::array::from_fn(|a| dims[a].max(patch[a]));
    let lo: [usize; 3] = std::array::from_fn(|a| (pdims[a] - dims[a]) / 2);
    let np: usize = pdims.iter().product();
    let mut padded = vec![0.0f32; np];
    for z in 0..dims[0] {
        for y in 0..dims[1] {
            let src = (z * dims[1] + y) * dims[2];
            let dst = ((z + lo[0]) * pdims[1] + y + lo[1]) * pdims[2] + lo[2];
            padded[dst..dst + dims[2]].copy_from_slice(&data[src..src + dims[2]]);
        }
    }

    let steps: [Vec<usize>; 3] = std::array::from_fn(|a| window_steps(pdims[a], patch[a], tile_step));
    let windows = steps.iter().map(Vec::len).product::<usize>();
    // One window covers everything: weight it uniformly (a Gaussian would cancel out anyway,
    // apart from rounding where it is ~1e-10 at the corners).
    let gauss = if windows == 1 { vec![1.0f32; patch.iter().product()] } else { gaussian_importance(patch) };
    let pp: usize = patch.iter().product();
    let shape = [1, 1, patch[0], patch[1], patch[2]];
    let plan = model.plan_for(&[&shape])?;
    let flips: &[[bool; 3]] = if mirror_tta {
        &[[false, false, false], [true, false, false], [false, true, false], [false, false, true],
          [true, true, false], [true, false, true], [false, true, true], [true, true, true]]
    } else {
        &[[false, false, false]]
    };
    let total = windows * flips.len();
    let mut done = 0;

    let mut acc = vec![0.0f32; channels * np];
    let mut weight = vec![0.0f32; np];
    let mut input = vec![0.0f32; pp];
    let mut pred = vec![0.0f32; channels * pp];
    for &s0 in &steps[0] {
        for &s1 in &steps[1] {
            for &s2 in &steps[2] {
                for_each_patch_row(pdims, patch, [s0, s1, s2], |src, dst| {
                    input[dst..dst + patch[2]].copy_from_slice(&padded[src..src + patch[2]]);
                });
                pred.iter_mut().for_each(|v| *v = 0.0);
                for flip in flips {
                    let x = flip3(&input, patch, *flip);
                    let out = plan.run_single(&Tensor::new(shape.to_vec(), x))?;
                    if out.shape != [1, channels, patch[0], patch[1], patch[2]] {
                        return Err(OnnxError::Run(format!("unexpected network output shape {:?}", out.shape)));
                    }
                    for c in 0..channels {
                        let back = flip3(&out.data[c * pp..(c + 1) * pp], patch, *flip);
                        pred[c * pp..(c + 1) * pp].iter_mut().zip(back).for_each(|(p, v)| *p += v);
                    }
                    done += 1;
                    progress(done, total);
                }
                let inv = 1.0 / flips.len() as f32;
                for_each_patch_row(pdims, patch, [s0, s1, s2], |dst, src| {
                    for k in 0..patch[2] {
                        let g = gauss[src + k];
                        weight[dst + k] += g;
                        for c in 0..channels {
                            acc[c * np + dst + k] += pred[c * pp + src + k] * inv * g;
                        }
                    }
                });
            }
        }
    }

    // Normalise and crop the padding back off.
    let n: usize = dims.iter().product();
    let mut logits = vec![0.0f32; channels * n];
    for c in 0..channels {
        for z in 0..dims[0] {
            for y in 0..dims[1] {
                let src = ((z + lo[0]) * pdims[1] + y + lo[1]) * pdims[2] + lo[2];
                let dst = (z * dims[1] + y) * dims[2];
                for x in 0..dims[2] {
                    logits[c * n + dst + x] = acc[c * np + src + x] / weight[src + x];
                }
            }
        }
    }
    Ok(logits)
}

/// Visit the rows of a patch at `origin` inside a volume: `f(volume_row_start, patch_row_start)`.
pub(super) fn for_each_patch_row(vdims: [usize; 3], patch: [usize; 3], origin: [usize; 3], mut f: impl FnMut(usize, usize)) {
    for z in 0..patch[0] {
        for y in 0..patch[1] {
            let v = ((origin[0] + z) * vdims[1] + origin[1] + y) * vdims[2] + origin[2];
            f(v, (z * patch[1] + y) * patch[2]);
        }
    }
}

/// Mirror a C-order patch along the flagged axes (an involution).
pub(super) fn flip3(src: &[f32], dims: [usize; 3], flip: [bool; 3]) -> Vec<f32> {
    if flip == [false; 3] {
        return src.to_vec();
    }
    let mut out = vec![0.0f32; src.len()];
    let m = |i: usize, a: usize| if flip[a] { dims[a] - 1 - i } else { i };
    for z in 0..dims[0] {
        for y in 0..dims[1] {
            for x in 0..dims[2] {
                out[(m(z, 0) * dims[1] + m(y, 1)) * dims[2] + m(x, 2)] = src[(z * dims[1] + y) * dims[2] + x];
            }
        }
    }
    out
}

/// Resample logits back to the cropped grid, argmax, un-crop. Returns the column-major mask.
fn postprocess(logits: &[f32], pre: &Preprocessed) -> Vec<u8> {
    let n = pre.data.len();
    let bbox = pre.bbox;
    let cdims = [bbox[0].1 - bbox[0].0, bbox[1].1 - bbox[1].0, bbox[2].1 - bbox[2].0];
    let back = |c: usize| -> Vec<f64> {
        let ch: Vec<f64> = logits[c * n..(c + 1) * n].iter().map(|&v| v as f64).collect();
        resample_nnunet(&ch, pre.dims, cdims, TARGET_SPACING, pre.spacing, 1)
    };
    let (bg, fg) = (back(0), back(1));
    let fd = pre.full_dims;
    let mut mask = vec![0u8; fd.iter().product()];
    let mut i = 0;
    for z in bbox[0].0..bbox[0].1 {
        for y in bbox[1].0..bbox[1].1 {
            let row = (z * fd[1] + y) * fd[2];
            for x in bbox[2].0..bbox[2].1 {
                // argmax over (background, brain); a tie goes to background like torch.argmax
                mask[row + x] = (fg[i] > bg[i]) as u8;
                i += 1;
            }
        }
    }
    mask
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn window_steps_match_nnunet() {
        // nnU-Net docstring example: image 110, patch 64, step 0.5 -> 0, 23, 46
        assert_eq!(window_steps(110, 64, 0.5), vec![0, 23, 46]);
        assert_eq!(window_steps(205, 96, 0.5), vec![0, 36, 73, 109]);
        assert_eq!(window_steps(192, 192, 0.5), vec![0]);
        assert_eq!(window_steps(205, 192, 1.0), vec![0, 13]);
    }

    #[test]
    fn full_volume_patch_rounds_up_to_the_divisor() {
        // (z, y, x): z to a multiple of 16, y and x to multiples of 32
        assert_eq!(full_volume_patch([150, 240, 240]), [160, 256, 256]);
        assert_eq!(full_volume_patch([160, 256, 224]), [160, 256, 224]);
        assert_eq!(full_volume_patch([1, 1, 1]), [16, 32, 32]);
        // ... so the sliding window over it is a single window
        let p = full_volume_patch([150, 240, 241]);
        assert!([150, 240, 241].iter().zip(p).all(|(&d, p)| window_steps(p.max(d), p, 1.0) == vec![0]));
    }

    #[test]
    fn default_is_full_volume() {
        assert_eq!(HdBetParams::default().inference, HdBetInference::FullVolume);
        assert!(!HdBetParams::default().mirror_tta);
        assert_eq!(HdBetParams::low_memory().inference, HdBetInference::LOW_MEMORY);
    }

    #[test]
    fn gaussian_peaks_at_centre_with_value_10() {
        let g = gaussian_importance([16, 32, 32]);
        let centre = (8 * 32 + 16) * 32 + 16;
        assert!((g[centre] - 10.0).abs() < 1e-6);
        assert!(g.iter().all(|&v| v > 0.0 && v <= 10.0));
    }

    #[test]
    fn separate_z_rule() {
        assert_eq!(separate_z_axis([4.0, 0.9, 0.9], [1.0; 3]), Some(0));
        assert_eq!(separate_z_axis([1.0; 3], [0.9, 0.9, 4.0]), Some(2));
        assert_eq!(separate_z_axis([1.2, 0.9, 0.9], [1.0; 3]), None);
        // two equally coarse axes -> no separate-z (nnU-Net: len(axis) == 2)
        assert_eq!(separate_z_axis([4.0, 4.0, 1.0], [1.0; 3]), None);
    }

    /// The whole step's peak against [`hd_bet_memory_bytes`]: native peaks of `hd_bet` on a
    /// synthetic head filling the field of view (`examples/inference_memory.rs hdbet`, counting
    /// allocator, single-threaded). The estimate must cover each and sit within 5 % of it.
    #[test]
    fn memory_estimate_covers_measured_peaks() {
        use crate::models::Heap;
        let window = HdBetParams { inference: HdBetInference::SLIDING_WINDOW, ..HdBetParams::default() };
        let cases = [
            ((96, 96, 64), (2.0, 2.0, 2.0), HdBetParams::default(), 4_326_663_752u64),
            ((184, 256, 128), (1.0, 1.0, 1.0), HdBetParams::default(), 5_675_605_064),
            ((256, 288, 48), (0.8, 0.8, 3.0), HdBetParams::default(), 7_352_436_552),
            ((240, 240, 160), (1.0, 1.0, 1.0), HdBetParams::default(), 9_257_042_005),
            ((160, 192, 96), (1.5, 1.5, 1.5), HdBetParams::default(), 9_375_081_688),
            ((184, 256, 128), (1.0, 1.0, 1.0), window.clone(), 3_371_351_512),
            ((320, 320, 240), (0.6, 0.6, 0.6), window, 3_352_739_280),
            ((184, 256, 128), (1.0, 1.0, 1.0), HdBetParams::low_memory(), 1_272_213_986),
            ((256, 288, 48), (0.8, 0.8, 3.0), HdBetParams::low_memory(), 1_287_419_762),
            ((320, 320, 240), (0.6, 0.6, 0.6), HdBetParams::low_memory(), 1_257_796_066),
        ];
        for (dims, voxel, params, peak) in cases {
            let grid = Grid::new(dims.0, dims.1, dims.2, voxel.0, voxel.1, voxel.2);
            let est = hd_bet_memory_bytes(&grid, &params, Heap::Native);
            let rel = (est as f64 - peak as f64) / peak as f64;
            assert!((0.0..=0.05).contains(&rel), "{dims:?} at {voxel:?}, {:?}: estimate {est} vs measured {peak}", params.inference);
            assert!(hd_bet_memory_bytes(&grid, &params, Heap::Wasm32) > est, "{dims:?}: the WASM heap needs more");
        }
    }

    #[test]
    fn network_input_is_the_field_of_view_at_1mm() {
        // 0.8×0.8×3 mm GRE, 256×288×48: 205×230×144 mm, padded to multiples of 32×32×16 (z, y, x).
        let grid = Grid::new(256, 288, 48, 0.8, 0.8, 3.0);
        assert_eq!(hd_bet_network_input(&grid, &HdBetParams::default()), [144, 256, 224]);
        assert_eq!(hd_bet_network_input(&grid, &HdBetParams::low_memory()), [64, 128, 128]);
    }

    #[test]
    fn flip_is_an_involution() {
        let d = [2, 3, 4];
        let v: Vec<f32> = (0..24).map(|i| i as f32).collect();
        let f = flip3(&v, d, [true, false, true]);
        assert_ne!(f, v);
        assert_eq!(flip3(&f, d, [true, false, true]), v);
        // voxel (0,0,0) of the flipped patch is source voxel (z=1, y=0, x=3)
        assert_eq!(f[0], v[(3 * 4) + 3]);
    }

    #[test]
    fn bbox_of_nonzero() {
        let mut d = vec![0.0; 3 * 4 * 5];
        let at = |z: usize, y: usize, x: usize| (z * 4 + y) * 5 + x;
        d[at(1, 2, 3)] = 1.0;
        d[at(2, 1, 1)] = -2.0;
        assert_eq!(nonzero_bbox(&d, [3, 4, 5]), Some([(1, 3), (1, 3), (1, 4)]));
        assert_eq!(nonzero_bbox(&[0.0; 8], [2, 2, 2]), None);
    }

    // ---- parity against nnU-Net intermediates (ref_hdbet.py in qsm-ci scripts/onnx-export) ----

    /// Minimal little-endian f32 C-order `.npy` reader.
    fn read_npy_f32(path: &str) -> (Vec<usize>, Vec<f32>) {
        let b = std::fs::read(path).unwrap_or_else(|e| panic!("{path}: {e}"));
        assert_eq!(&b[..6], b"\x93NUMPY");
        let (hlen, off) = if b[6] == 1 { (u16::from_le_bytes([b[8], b[9]]) as usize, 10) } else {
            (u32::from_le_bytes([b[8], b[9], b[10], b[11]]) as usize, 12)
        };
        let header = std::str::from_utf8(&b[off..off + hlen]).unwrap();
        assert!(header.contains("'<f4'") && header.contains("'fortran_order': False"), "{header}");
        let shape_str = header.split("'shape': (").nth(1).unwrap().split(')').next().unwrap();
        let shape: Vec<usize> = shape_str.split(',').filter_map(|s| s.trim().parse().ok()).collect();
        let data = b[off + hlen..].as_chunks::<4>().0.iter().map(|c| f32::from_le_bytes(*c)).collect();
        (shape, data)
    }

    fn ref_dir() -> Option<String> {
        std::env::var("HDBET_REF_DIR").ok()
    }

    fn load_case(dir: &str, case: &str) -> (Vec<f64>, Grid) {
        let nii = crate::io::read_nifti_file(std::path::Path::new(&format!("{dir}/{case}_mag.nii.gz"))).unwrap();
        let (nx, ny, nz) = nii.dims;
        let (vx, vy, vz) = nii.voxel_size;
        (nii.data, Grid::new(nx, ny, nz, vx, vy, vz))
    }

    /// Crop + z-score + resample vs nnU-Net's `run_case_npy` for the three reference cases.
    #[test]
    #[ignore]
    fn preprocessing_matches_nnunet() {
        let Some(dir) = ref_dir() else { return eprintln!("HDBET_REF_DIR not set; skipping") };
        for case in ["A", "B", "C"] {
            let (mag, grid) = load_case(&dir, case);
            let pre = preprocess(&mag, &grid).unwrap();
            let (shape, want) = read_npy_f32(&format!("{dir}/{case}_preprocessed.npy"));
            assert_eq!(pre.dims.to_vec(), shape, "case {case}: shape");
            let err = pre.data.iter().zip(&want).map(|(a, b)| (a - b).abs()).fold(0.0f32, f32::max);
            eprintln!("case {case}: dims {:?}, bbox {:?}, max |d| {err:e}", pre.dims, pre.bbox);
            assert!(err < 1e-3, "case {case}: preprocessed max |d| = {err}");
        }
    }

    /// Resample-back + argmax + un-crop applied to nnU-Net's own logits reproduces the CLI mask.
    #[test]
    #[ignore]
    fn postprocessing_matches_nnunet() {
        let Some(dir) = ref_dir() else { return eprintln!("HDBET_REF_DIR not set; skipping") };
        for case in ["A", "B", "C"] {
            let (mag, grid) = load_case(&dir, case);
            let pre = preprocess(&mag, &grid).unwrap();
            let (_, logits) = read_npy_f32(&format!("{dir}/{case}_logits.npy"));
            let mask = postprocess(&logits, &pre);
            let cli = crate::io::read_nifti_file(std::path::Path::new(&format!("{dir}/{case}_mask_cli.nii.gz"))).unwrap();
            let diff = mask.iter().zip(&cli.data).filter(|(&m, &c)| (m != 0) != (c > 0.5)).count();
            eprintln!("case {case}: {diff} voxels differ from the CLI mask");
            assert!(diff * 100_000 < mask.len(), "case {case}: {diff} differing voxels");
        }
    }
}