qsm_core/separation/
susep_net.rs

1//! SUSEP-Net deep-learning χ-separation (`onnx` feature).
2//!
3//! SUSEP-Net (Li/Gao/Sun 2025) is a dual-branch 3D U-Net that maps three guidance
4//! maps — QSM (χ_total, ppm), R2′ (Hz), local field (ppm) — to paramagnetic (χ+)
5//! and diamagnetic (χ−) source magnitudes. Clean NCDHW ONNX export (three inputs
6//! `qsm`,`r2prime`,`lfs`; two outputs `chi_pos`,`chi_neg`).
7//!
8//! Recipe (mirrors the authors' `recon.py`): z-score each input by the training
9//! stats, zero outside the mask, post-pad each dim to a multiple of 8, run,
10//! de-normalize the outputs, crop, and mask. The network's ReLU makes both
11//! outputs non-negative magnitudes; we return χ− as a signed (≤ 0) value to match
12//! the crate's separation convention `(chi_pos ≥ 0, chi_neg ≤ 0, chi_total)`.
13//!
14//! Weights are not bundled; the caller passes the exported `susep-net.onnx` bytes
15//! (see [`crate::models`]).
16
17use crate::grid::Grid;
18use crate::models::onnx::{OnnxModel, OnnxError, Tensor};
19
20/// Training z-score constants for SUSEP-Net (`all_mean_std.mat`): each field is
21/// `(mean, std)`. Inputs are normalized `(x-mean)/std`; outputs de-normalized
22/// `y*std + mean`.
23#[derive(Clone, Copy, Debug)]
24pub struct SusepNetNorm {
25    pub qsm: (f64, f64),
26    pub lfs: (f64, f64),
27    pub r2prime: (f64, f64),
28    pub chi_pos: (f64, f64),
29    pub chi_neg: (f64, f64),
30}
31
32impl Default for SusepNetNorm {
33    /// Constants shipped with the released `SUSEPNet.pth`.
34    fn default() -> Self {
35        Self {
36            qsm: (-6.0663105e-05, 0.023533047),
37            lfs: (-4.8702253e-05, 0.012554166),
38            r2prime: (4.7629275, 10.889079),
39            chi_pos: (0.0089528897, 0.025519046),
40            chi_neg: (0.0090135528, 0.019603666),
41        }
42    }
43}
44
45/// Run SUSEP-Net χ-separation.
46///
47/// * `local_field_ppm`, `qsm` (χ_total, ppm), `r2prime` (Hz) — column-major `(nx,ny,nz)`.
48/// * `mask` — binary brain mask (same layout).
49/// * `model_onnx` — bytes of the exported `susep-net.onnx`.
50/// * `norm` — training normalization constants.
51///
52/// Returns `(chi_pos ≥ 0, chi_neg ≤ 0, chi_total = chi_pos + chi_neg)` in ppm,
53/// masked, in the same layout.
54#[allow(clippy::too_many_arguments)]
55pub fn susep_net(
56    local_field_ppm: &[f64],
57    qsm: &[f64],
58    r2prime: &[f64],
59    mask: &[u8],
60    grid: &Grid,
61    model_onnx: &[u8],
62    norm: &SusepNetNorm,
63) -> Result<(Vec<f64>, Vec<f64>, Vec<f64>), OnnxError> {
64    let (nx, ny, nz) = grid.dims;
65    let n = nx * ny * nz;
66    for (name, v) in [("field", local_field_ppm), ("qsm", qsm), ("r2prime", r2prime)] {
67        assert_eq!(v.len(), n, "{name} length must match grid");
68    }
69    assert_eq!(mask.len(), n, "mask length must match grid");
70
71    // Post-pad each dim to a multiple of 8 (three 2× pools).
72    let (px, py, pz) = (nx.div_ceil(8) * 8, ny.div_ceil(8) * 8, nz.div_ceil(8) * 8);
73
74    // z-score + mask + repack column-major (nx,ny,nz) → row-major NCDHW [1,1,px,py,pz].
75    let pack = |src: &[f64], (mean, std): (f64, f64)| -> Tensor {
76        let inv = 1.0 / std;
77        let mut buf = vec![0.0f32; px * py * pz];
78        for z in 0..nz {
79            for y in 0..ny {
80                for x in 0..nx {
81                    let i = x + nx * (y + ny * z);
82                    if mask[i] != 0 {
83                        let dst = (x * py + y) * pz + z;
84                        buf[dst] = ((src[i] - mean) * inv) as f32;
85                    }
86                }
87            }
88        }
89        Tensor::new(vec![1, 1, px, py, pz], buf)
90    };
91
92    // Input order must match the exported graph: qsm, r2prime, lfs.
93    let inputs = [
94        pack(qsm, norm.qsm),
95        pack(r2prime, norm.r2prime),
96        pack(local_field_ppm, norm.lfs),
97    ];
98    let model = OnnxModel::load(model_onnx)?;
99    let outs = model.run(&inputs)?;
100    if outs.len() < 2 {
101        return Err(OnnxError::Run(format!("expected 2 outputs, got {}", outs.len())));
102    }
103
104    // De-normalize, crop, mask, unpack. χ− is a magnitude → return signed (≤ 0).
105    let mut chi_pos = vec![0.0f64; n];
106    let mut chi_neg = vec![0.0f64; n];
107    let mut chi_total = vec![0.0f64; n];
108    let (pm, ps) = norm.chi_pos;
109    let (nm, ns) = norm.chi_neg;
110    for z in 0..nz {
111        for y in 0..ny {
112            for x in 0..nx {
113                let i = x + nx * (y + ny * z);
114                if mask[i] != 0 {
115                    let src = (x * py + y) * pz + z;
116                    let pos = outs[0].data[src] as f64 * ps + pm;
117                    let neg_mag = outs[1].data[src] as f64 * ns + nm;
118                    chi_pos[i] = pos;
119                    chi_neg[i] = -neg_mag;
120                    chi_total[i] = pos - neg_mag;
121                }
122            }
123        }
124    }
125    Ok((chi_pos, chi_neg, chi_total))
126}
qsm_core/separation/susep_net.rs

qsm_core/separation/
susep_net.rs
