qsm_core/inversion/
qsmnet.rs

1//! QSMnet deep-learning dipole inversion (`onnx` feature).
2//!
3//! QSMnet is a 3D U-Net (SNU-LIST) that maps a local field (ppm) to
4//! susceptibility (ppm). The upstream weights are TensorFlow, but we ship a
5//! *clean* PyTorch re-export (`scripts/onnx-export/export_qsmnet.py`): the plain
6//! U-Net rebuilt in PyTorch with the TF weights ported in, giving a tidy NCDHW
7//! ONNX (`[1, 1, X, Y, Z]`) that the pure-Rust `tract` engine runs — unlike the
8//! `tf2onnx` graph, whose NHWC↔NCHW Reshape/Transpose ops tract can't analyse.
9//! This mirrors the authors' inference (`Code/inference.py`): normalize by the
10//! dataset mean/std shipped with the checkpoint, centered zero-pad each dim to a
11//! multiple of 16 (four pool/deconv levels), run, crop, de-normalize, and mask.
12//!
13//! Weights are not bundled; the caller passes the exported `qsmnet.onnx` bytes
14//! (see [`crate::models`]).
15
16use crate::grid::Grid;
17use crate::models::onnx::{OnnxModel, OnnxError, Tensor};
18
19/// Dataset normalization constants stored beside a QSMnet checkpoint
20/// (`norm_factor_<name>.mat`): `field_n = (field - in_mean)/in_std`, and
21/// `chi = out_std*pred + out_mean`.
22#[derive(Clone, Copy, Debug)]
23pub struct QsmnetNorm {
24    pub in_mean: f64,
25    pub in_std: f64,
26    pub out_mean: f64,
27    pub out_std: f64,
28}
29
30impl QsmnetNorm {
31    /// Constants for the `QSMnet_64` checkpoint.
32    pub fn qsmnet() -> Self {
33        Self { in_mean: 0.0, in_std: 0.01, out_mean: 0.0, out_std: 0.0317 }
34    }
35
36    /// Constants for the `QSMnet+_64` checkpoint.
37    pub fn qsmnet_plus() -> Self {
38        Self { in_mean: 0.0, in_std: 0.0205, out_mean: 0.0, out_std: 0.0734 }
39    }
40}
41
42impl Default for QsmnetNorm {
43    /// Constants for the `QSMnet_64` checkpoint.
44    fn default() -> Self {
45        Self::qsmnet()
46    }
47}
48
49/// Run QSMnet dipole inversion.
50///
51/// * `local_field_ppm` — local tissue field (ppm), column-major `(nx,ny,nz)`.
52/// * `mask` — binary brain mask (same layout); applied to the output.
53/// * `model_onnx` — bytes of the exported `qsmnet.onnx` (NDHWC input/output).
54/// * `norm` — dataset normalization constants for this checkpoint.
55///
56/// Returns susceptibility (ppm), masked, in the same layout.
57pub fn qsmnet(
58    local_field_ppm: &[f64],
59    mask: &[u8],
60    grid: &Grid,
61    model_onnx: &[u8],
62    norm: &QsmnetNorm,
63) -> Result<Vec<f64>, OnnxError> {
64    let (nx, ny, nz) = grid.dims;
65    let n = nx * ny * nz;
66    assert_eq!(local_field_ppm.len(), n, "field length must match grid");
67    assert_eq!(mask.len(), n, "mask length must match grid");
68
69    // Centered padding to a multiple of 16 (matches the Python `pad_to_multiple`).
70    let pad = |s: usize| -> (usize, usize) {
71        let total = (16 - s % 16) % 16;
72        (total / 2, s + total)
73    };
74    let (bx, px) = pad(nx);
75    let (by, py) = pad(ny);
76    let (bz, pz) = pad(nz);
77
78    // Normalize + repack column-major (nx,ny,nz) f64 → row-major NCDHW
79    // [1,1,px,py,pz] f32, centered.
80    let inv_std = 1.0 / norm.in_std;
81    let mut input = vec![0.0f32; px * py * pz];
82    for z in 0..nz {
83        for y in 0..ny {
84            for x in 0..nx {
85                let src = x + nx * (y + ny * z);
86                let dst = ((x + bx) * py + (y + by)) * pz + (z + bz);
87                input[dst] = ((local_field_ppm[src] - norm.in_mean) * inv_std) as f32;
88            }
89        }
90    }
91
92    let model = OnnxModel::load(model_onnx)?;
93    let out = model.run_single(&Tensor::new(vec![1, 1, px, py, pz], input))?;
94    let expected = [1, 1, px, py, pz];
95    if out.shape != expected {
96        return Err(OnnxError::Run(format!(
97            "unexpected output shape {:?}, expected {expected:?}",
98            out.shape
99        )));
100    }
101
102    // De-normalize, crop (centered), mask, unpack → column-major.
103    let mut chi = vec![0.0f64; n];
104    for z in 0..nz {
105        for y in 0..ny {
106            for x in 0..nx {
107                let dst = x + nx * (y + ny * z);
108                if mask[dst] != 0 {
109                    let src = ((x + bx) * py + (y + by)) * pz + (z + bz);
110                    chi[dst] = norm.out_std * out.data[src] as f64 + norm.out_mean;
111                }
112            }
113        }
114    }
115    Ok(chi)
116}
117
118/// Memory-bounded QSMnet via overlap-tiling — the fully-convolutional U-Net run patch-by-patch
119/// (for 32-bit WASM, where whole-volume [`qsmnet`] overflows the heap on clinical data). The
120/// dataset normalization is applied per value inside the tile loop and the net's pool depth
121/// requires a `/16` patch. Approximates whole-volume up to tile-boundary error; see
122/// [`crate::inversion::tiled`].
123pub fn qsmnet_tiled(
124    local_field_ppm: &[f64],
125    mask: &[u8],
126    grid: &Grid,
127    model_onnx: &[u8],
128    norm: &QsmnetNorm,
129    cfg: &super::tiled::TileConfig,
130    progress: impl FnMut(usize, usize),
131) -> Result<Vec<f64>, OnnxError> {
132    let model = OnnxModel::load(model_onnx)?;
133    let (in_mean, inv_std, out_std, out_mean) =
134        (norm.in_mean, 1.0 / norm.in_std, norm.out_std, norm.out_mean);
135    super::tiled::tiled_field_inversion(
136        local_field_ppm, mask, grid, &model, 16, cfg,
137        move |v| ((v - in_mean) * inv_std) as f32,
138        move |o| out_std * o as f64 + out_mean,
139        progress,
140    )
141}
qsm_core/inversion/qsmnet.rs

qsm_core/inversion/
qsmnet.rs
