qsm_core/inversion/
xqsm.rs

1//! xQSM deep-learning dipole inversion (`onnx` feature).
2//!
3//! xQSM is an octave-convolution U-Net (with a global residual) that maps a
4//! local tissue field (ppm) to susceptibility (ppm). It is orientation-agnostic
5//! and takes no dipole kernel — the field goes in, χ comes out. This mirrors the
6//! authors' pure-Python inference (`sunhongfu/xQSM`): centered zero-pad each
7//! dimension to a multiple of 8, run the net, crop back, and multiply the output
8//! by the mask. No normalization.
9//!
10//! Weights are not bundled; the caller passes the exported `xqsm.onnx` bytes
11//! (see [`crate::models`]).
12
13use crate::grid::Grid;
14use crate::models::onnx::{OnnxModel, OnnxError, Tensor};
15
16/// Run xQSM dipole inversion.
17///
18/// * `local_field_ppm` — local tissue field (ppm), column-major `(nx,ny,nz)`.
19/// * `mask` — binary brain mask (same layout); applied to the output.
20/// * `model_onnx` — bytes of the exported `xqsm.onnx`.
21///
22/// Returns susceptibility (ppm), masked, in the same layout.
23pub fn xqsm(
24    local_field_ppm: &[f64],
25    mask: &[u8],
26    grid: &Grid,
27    model_onnx: &[u8],
28) -> Result<Vec<f64>, OnnxError> {
29    let (nx, ny, nz) = grid.dims;
30    let n = nx * ny * nz;
31    assert_eq!(local_field_ppm.len(), n, "field length must match grid");
32    assert_eq!(mask.len(), n, "mask length must match grid");
33
34    // Centered padding offsets so each dim is a multiple of 8 (matches the
35    // Python reference's `_zero_pad`).
36    let pad = |s: usize| -> (usize, usize) {
37        let total = (8 - s % 8) % 8;
38        let before = total / 2;
39        (before, s + total) // (offset, padded size)
40    };
41    let (bx, px) = pad(nx);
42    let (by, py) = pad(ny);
43    let (bz, pz) = pad(nz);
44
45    // Repack column-major (nx,ny,nz) f64 → row-major (px,py,pz) f32, centered.
46    // The field is NOT masked before inference (unlike BFRnet).
47    let mut input = vec![0.0f32; px * py * pz];
48    for z in 0..nz {
49        for y in 0..ny {
50            for x in 0..nx {
51                let src = x + nx * (y + ny * z);
52                let dst = (z + bz) + pz * ((y + by) + py * (x + bx));
53                input[dst] = local_field_ppm[src] as f32;
54            }
55        }
56    }
57
58    let model = OnnxModel::load(model_onnx)?;
59    let out = model.run_single(&Tensor::new(vec![1, 1, px, py, pz], input))?;
60    if out.shape != [1, 1, px, py, pz] {
61        return Err(OnnxError::Run(format!(
62            "unexpected output shape {:?}, expected [1,1,{px},{py},{pz}]",
63            out.shape
64        )));
65    }
66
67    // Crop back (centered), mask, and unpack row-major → column-major.
68    let mut chi = vec![0.0f64; n];
69    for z in 0..nz {
70        for y in 0..ny {
71            for x in 0..nx {
72                let dst = x + nx * (y + ny * z);
73                if mask[dst] != 0 {
74                    let src = (z + bz) + pz * ((y + by) + py * (x + bx));
75                    chi[dst] = out.data[src] as f64;
76                }
77            }
78        }
79    }
80    Ok(chi)
81}
82
83/// Memory-bounded xQSM via overlap-tiling — the same fully-convolutional net run patch-by-
84/// patch so peak memory is bounded by one patch (for 32-bit WASM, where the whole-volume
85/// [`xqsm`] overflows the 4 GB heap on clinical-size data). Results approximate whole-volume
86/// xQSM up to tile-boundary error; see [`crate::inversion::tiled`].
87pub fn xqsm_tiled(
88    local_field_ppm: &[f64],
89    mask: &[u8],
90    grid: &Grid,
91    model_onnx: &[u8],
92    cfg: &super::tiled::TileConfig,
93    progress: impl FnMut(usize, usize),
94) -> Result<Vec<f64>, OnnxError> {
95    let model = OnnxModel::load(model_onnx)?;
96    // xQSM feeds the raw field (no normalization) and reads χ directly; pad each patch to /8.
97    super::tiled::tiled_field_inversion(
98        local_field_ppm, mask, grid, &model, 8, cfg, |v| v as f32, |o| o as f64, progress,
99    )
100}
qsm_core/inversion/xqsm.rs

qsm_core/inversion/
xqsm.rs
