qsm_core/inversion/
qsmgan.rs

1//! QSMGAN dipole inversion (`onnx` feature).
2//!
3//! QSMGAN (Chen 2020) inverts the dipole with a 3D U-Net **generator** refined by a
4//! Wasserstein GAN. Inference is patch-based with an increased receptive field: a
5//! 64³ local-field patch maps to a 48³ susceptibility patch (a centre crop baked
6//! into the exported graph — the "i64o48" scheme). The network itself is the only
7//! part in ONNX; this glue reproduces the authors' `recon.py` pipeline:
8//!
9//! * sign flip — the QSM-CI runner negates the local field before the net;
10//! * tile by non-overlapping 48³ **output** patches, each fed a 64³ **input** patch
11//!   (same centre, 8-voxel context margin, zero-padded at the volume edges);
12//! * `input_scale = 100`; output head `tanh`, so `χ = atanh(clip(out)) / 10`;
13//! * mask the result.
14//!
15//! The input is the SEPIA-style **local field in ppm** (already background-removed).
16//! Weights are not bundled; the caller passes the exported `qsmgan.onnx` bytes.
17
18use crate::grid::Grid;
19use crate::models::onnx::{OnnxError, OnnxModel, Tensor};
20
21const IPS: usize = 64; // input patch size
22const OPS: usize = 48; // output patch size
23const INPUT_SCALE: f64 = 100.0;
24const OUTPUT_SCALE: f64 = 10.0;
25
26/// Run QSMGAN on a background-removed local field (ppm), column-major `(nx,ny,nz)`.
27/// Returns susceptibility (ppm), masked, in the same layout.
28pub fn qsmgan(
29    local_field_ppm: &[f64],
30    mask: &[u8],
31    grid: &Grid,
32    model_onnx: &[u8],
33) -> Result<Vec<f64>, OnnxError> {
34    let (nx, ny, nz) = grid.dims;
35    let n = nx * ny * nz;
36    assert_eq!(local_field_ppm.len(), n, "field length must match grid");
37    assert_eq!(mask.len(), n, "mask length must match grid");
38
39    let model = OnnxModel::load(model_onnx)?;
40    let (half_i, half_o) = (IPS as i64 / 2, OPS as i64 / 2); // 32, 24
41    let (nxi, nyi, nzi) = (nx as i64, ny as i64, nz as i64);
42    // Negated local field (the runner's sign flip), zero outside the volume.
43    let field = |x: i64, y: i64, z: i64| -> f32 {
44        if x >= 0 && x < nxi && y >= 0 && y < nyi && z >= 0 && z < nzi {
45            (-INPUT_SCALE * local_field_ppm[x as usize + nx * (y as usize + ny * z as usize)]) as f32
46        } else {
47            0.0
48        }
49    };
50
51    let mut predict = vec![0.0f64; n];
52    // Output-patch centres: range(half_o, dim + half_o + 1, OPS) per axis.
53    let mut cx = half_o;
54    while cx <= nxi + half_o {
55        let mut cy = half_o;
56        while cy <= nyi + half_o {
57            let mut cz = half_o;
58            while cz <= nzi + half_o {
59                // In-volume extent of this 48³ output patch.
60                let (x0, y0, z0) = (cx - half_o, cy - half_o, cz - half_o); // ≥ 0
61                let xe = (x0 + OPS as i64).min(nxi);
62                let ye = (y0 + OPS as i64).min(nyi);
63                let ze = (z0 + OPS as i64).min(nzi);
64                let (px, py, pz) = (xe - x0, ye - y0, ze - z0);
65                if px <= 0 || py <= 0 || pz <= 0 {
66                    cz += OPS as i64;
67                    continue;
68                }
69
70                // 64³ input patch centred at (cx,cy,cz), row-major NCDHW.
71                let mut inp = vec![0.0f32; IPS * IPS * IPS];
72                for i in 0..IPS {
73                    for j in 0..IPS {
74                        for k in 0..IPS {
75                            let v = field(
76                                cx - half_i + i as i64,
77                                cy - half_i + j as i64,
78                                cz - half_i + k as i64,
79                            );
80                            inp[(i * IPS + j) * IPS + k] = v;
81                        }
82                    }
83                }
84                let out = model.run_single(&Tensor::new(vec![1, 1, IPS, IPS, IPS], inp))?;
85
86                // Place the in-bounds part of the 48³ output patch: χ = atanh(clip)/10.
87                for oi in 0..px as usize {
88                    for oj in 0..py as usize {
89                        for ok in 0..pz as usize {
90                            let o = (out.data[(oi * OPS + oj) * OPS + ok] as f64)
91                                .clamp(-0.999999, 0.999999);
92                            let dst = (x0 as usize + oi)
93                                + nx * ((y0 as usize + oj) + ny * (z0 as usize + ok));
94                            predict[dst] = o.atanh() / OUTPUT_SCALE;
95                        }
96                    }
97                }
98                cz += OPS as i64;
99            }
100            cy += OPS as i64;
101        }
102        cx += OPS as i64;
103    }
104
105    for i in 0..n {
106        if mask[i] == 0 {
107            predict[i] = 0.0;
108        }
109    }
110    Ok(predict)
111}
qsm_core/inversion/qsmgan.rs

qsm_core/inversion/
qsmgan.rs
