qsm_core/inversion/
iqsm_plus.rs

1//! iQSM+ single-step deep-learning reconstruction (`onnx` feature).
2//!
3//! iQSM+ (Gao 2024) extends [iQSM](crate::inversion::iqsm) with orientation-adaptive
4//! latent feature editing (OA-LFE): the B0 direction (`z_prjs`) is a network input,
5//! so oblique/sagittal/coronal acquisitions reconstruct correctly. The exported
6//! graph takes six inputs: `phase`, `mask`, `te` (s), `b0` (T), `z_prjs` (B0 dir,
7//! `[1,1,3]`), and `border` (the LoT boundary mask, as for iQSM).
8//!
9//! Pipeline (mirrors the authors' `inference.run_iqsm_plus`): flip the phase sign,
10//! erode the mask by a radius-3 sphere, **crop to the brain bounding box + 16-voxel
11//! margin**, centre-pad to a multiple of 16, run, ×mask, then paste the result back
12//! into the full grid. Multi-echo data is combined with magnitude·TE² weighting.
13//!
14//! Only axial-ish acquisitions are handled directly here; the authors' extra
15//! axis-permutation for strongly oblique fields (`|dir_y| > |dir_z|`) is not applied.
16//!
17//! Weights are not bundled; the caller passes the exported `iqsm-plus.onnx` bytes.
18
19use crate::grid::Grid;
20use crate::inversion::iqsm::sphere_erode;
21use crate::models::onnx::{OnnxModel, OnnxError, Tensor};
22
23/// Run iQSM+ on a single echo of wrapped phase.
24///
25/// * `phase_rad`, `mask` — column-major `(nx,ny,nz)`.
26/// * `te` (s), `b0` (T), `b0_dir` — acquisition parameters (`b0_dir` normalized internally).
27/// * `phase_sign` (`-1` default), `eroded_rad` (`3` default).
28///
29/// Returns susceptibility (ppm), masked, in the same layout.
30#[allow(clippy::too_many_arguments)]
31pub fn iqsm_plus(
32    phase_rad: &[f64],
33    mask: &[u8],
34    grid: &Grid,
35    te: f64,
36    b0: f64,
37    b0_dir: (f64, f64, f64),
38    phase_sign: f64,
39    eroded_rad: i32,
40    model_onnx: &[u8],
41) -> Result<Vec<f64>, OnnxError> {
42    let model = OnnxModel::load(model_onnx)?;
43    iqsm_plus_with(&model, phase_rad, mask, grid, te, b0, b0_dir, phase_sign, eroded_rad)
44}
45
46/// Multi-echo iQSM+: reconstruct each echo and combine with magnitude·TE² weights.
47#[allow(clippy::too_many_arguments)]
48pub fn iqsm_plus_multi_echo(
49    phases: &[&[f64]],
50    magnitudes: &[&[f64]],
51    mask: &[u8],
52    grid: &Grid,
53    tes: &[f64],
54    b0: f64,
55    b0_dir: (f64, f64, f64),
56    phase_sign: f64,
57    eroded_rad: i32,
58    model_onnx: &[u8],
59) -> Result<Vec<f64>, OnnxError> {
60    assert_eq!(phases.len(), tes.len(), "one TE per echo");
61    let n = grid.n_total();
62    let model = OnnxModel::load(model_onnx)?;
63    let mut acc = vec![0.0f64; n];
64    let mut wsum = vec![0.0f64; n];
65    for (e, &phase) in phases.iter().enumerate() {
66        let chi = iqsm_plus_with(&model, phase, mask, grid, tes[e], b0, b0_dir, phase_sign, eroded_rad)?;
67        let te2 = tes[e] * tes[e];
68        for i in 0..n {
69            let w = magnitudes.get(e).map(|m| m[i]).unwrap_or(1.0) * te2;
70            acc[i] += w * chi[i];
71            wsum[i] += w;
72        }
73    }
74    for i in 0..n {
75        acc[i] = if wsum[i] > 0.0 { acc[i] / wsum[i] } else { 0.0 };
76    }
77    Ok(acc)
78}
79
80#[allow(clippy::too_many_arguments)]
81fn iqsm_plus_with(
82    model: &OnnxModel,
83    phase_rad: &[f64],
84    mask: &[u8],
85    grid: &Grid,
86    te: f64,
87    b0: f64,
88    b0_dir: (f64, f64, f64),
89    phase_sign: f64,
90    eroded_rad: i32,
91) -> Result<Vec<f64>, OnnxError> {
92    let (nx, ny, nz) = grid.dims;
93    let n = nx * ny * nz;
94    assert_eq!(phase_rad.len(), n, "phase length must match grid");
95    assert_eq!(mask.len(), n, "mask length must match grid");
96    let idx = |x: usize, y: usize, z: usize| x + nx * (y + ny * z);
97
98    let eroded = if eroded_rad > 0 {
99        sphere_erode(mask, grid, eroded_rad)
100    } else {
101        mask.to_vec()
102    };
103
104    // Brain bounding box of the eroded mask, expanded by 16 and clamped.
105    let (mut x0, mut y0, mut z0) = (nx, ny, nz);
106    let (mut x1, mut y1, mut z1) = (0usize, 0usize, 0usize);
107    let mut any = false;
108    for z in 0..nz {
109        for y in 0..ny {
110            for x in 0..nx {
111                if eroded[idx(x, y, z)] != 0 {
112                    any = true;
113                    x0 = x0.min(x); y0 = y0.min(y); z0 = z0.min(z);
114                    x1 = x1.max(x); y1 = y1.max(y); z1 = z1.max(z);
115                }
116            }
117        }
118    }
119    if !any {
120        return Ok(vec![0.0; n]);
121    }
122    let pad = 16usize;
123    let bx0 = x0.saturating_sub(pad);
124    let by0 = y0.saturating_sub(pad);
125    let bz0 = z0.saturating_sub(pad);
126    let bx1 = (x1 + 1 + pad).min(nx);
127    let by1 = (y1 + 1 + pad).min(ny);
128    let bz1 = (z1 + 1 + pad).min(nz);
129    let (cx, cy, cz) = (bx1 - bx0, by1 - by0, bz1 - bz0);
130
131    // Centre-pad the cropped box to a multiple of 16.
132    let mpad = |s: usize| -> (usize, usize) {
133        let total = (16 - s % 16) % 16;
134        (total / 2, s + total)
135    };
136    let (ox, px) = mpad(cx);
137    let (oy, py) = mpad(cy);
138    let (oz, pz) = mpad(cz);
139
140    // Fill NCDHW row-major [1,1,px,py,pz] phase/mask/border from the cropped box.
141    let mut phase_t = vec![0.0f32; px * py * pz];
142    let mut mask_t = vec![0.0f32; px * py * pz];
143    let mut border = vec![1.0f32; px * py * pz];
144    let rdst = |i: usize, j: usize, k: usize| (k + oz) + pz * ((j + oy) + py * (i + ox));
145    for k in 0..cz {
146        for j in 0..cy {
147            for i in 0..cx {
148                let src = idx(bx0 + i, by0 + j, bz0 + k);
149                let dst = rdst(i, j, k);
150                phase_t[dst] = (phase_sign * phase_rad[src]) as f32;
151                mask_t[dst] = eroded[src] as f32;
152            }
153        }
154    }
155    for a in 0..px {
156        for b in 0..py {
157            for c in 0..pz {
158                if a == 0 || a == px - 1 || b == 0 || b == py - 1 || c == 0 || c == pz - 1 {
159                    border[c + pz * (b + py * a)] = 0.0;
160                }
161            }
162        }
163    }
164
165    let norm = (b0_dir.0.powi(2) + b0_dir.1.powi(2) + b0_dir.2.powi(2)).sqrt();
166    let zdir = if norm > 0.0 { norm } else { 1.0 };
167    let z_prjs = vec![
168        (b0_dir.0 / zdir) as f32,
169        (b0_dir.1 / zdir) as f32,
170        (b0_dir.2 / zdir) as f32,
171    ];
172
173    let shape = vec![1, 1, px, py, pz];
174    let inputs = [
175        Tensor::new(shape.clone(), phase_t),
176        Tensor::new(shape.clone(), mask_t),
177        Tensor::new(vec![1], vec![te as f32]),
178        Tensor::new(vec![1], vec![b0 as f32]),
179        Tensor::new(vec![1, 1, 3], z_prjs),
180        Tensor::new(shape, border),
181    ];
182    let out = model.run(&inputs)?;
183    let chi_pad = &out[0].data;
184
185    // ×mask, crop pad, paste back into the full grid at the bounding box.
186    let mut chi = vec![0.0f64; n];
187    for k in 0..cz {
188        for j in 0..cy {
189            for i in 0..cx {
190                let src = idx(bx0 + i, by0 + j, bz0 + k);
191                if eroded[src] != 0 {
192                    chi[src] = chi_pad[rdst(i, j, k)] as f64;
193                }
194            }
195        }
196    }
197    Ok(chi)
198}
qsm_core/inversion/iqsm_plus.rs

qsm_core/inversion/
iqsm_plus.rs
