qsm_core/inversion/
autoqsm.rs

1//! AutoQSM single-step deep-learning reconstruction (`onnx` feature).
2//!
3//! AutoQSM (Wei 2019) is a patch V-Net that maps a **total** field (ppm) directly
4//! to susceptibility (ppm) — no brain extraction, no separate background removal.
5//! The network takes a fixed 64³ input patch and returns the central 32³. Whole
6//! volumes are reconstructed by overlap-tiled sliding-window inference: 32³ output
7//! patches at stride 24 (8-voxel overlap), each from a 64³ input patch with a
8//! 16-voxel context margin, edge-padded and linearly blended across the overlap.
9//! This mirrors the authors' `util.data_predict` / `patch_process`.
10//!
11//! Upstream weights are Keras; QSM.rs ships a clean PyTorch re-export
12//! (`scripts/onnx-export/export_autoqsm.py`) so `tract` runs it. The network is
13//! all Conv3D+ReLU with no normalization, so no input scaling is applied.
14//!
15//! Weights are not bundled; the caller passes the exported `autoqsm.onnx` bytes.
16
17use crate::grid::Grid;
18use crate::models::onnx::{OnnxModel, OnnxError, Tensor};
19
20const IN: usize = 64;
21const OUT: usize = 32;
22const MARGIN: usize = (IN - OUT) / 2; // 16
23const SHIFT: usize = 24;
24const OVERLAP: usize = OUT - SHIFT; // 8
25
26#[inline]
27fn ri(x: usize, y: usize, z: usize, dy: usize, dz: usize) -> usize {
28    (x * dy + y) * dz + z // row-major (X,Y,Z), matches numpy C-order
29}
30
31/// Run AutoQSM on a total field map.
32///
33/// * `field_ppm` — total field (ppm), column-major `(nx,ny,nz)`.
34/// * `mask` — binary brain mask; applied to the output (AutoQSM has no brain
35///   extraction, so the raw output is whole-head).
36/// * `model_onnx` — bytes of the exported `autoqsm.onnx` (fixed 64³→32³).
37///
38/// Returns susceptibility (ppm), masked, in the same layout.
39pub fn autoqsm(
40    field_ppm: &[f64],
41    mask: &[u8],
42    grid: &Grid,
43    model_onnx: &[u8],
44) -> Result<Vec<f64>, OnnxError> {
45    let (nx, ny, nz) = grid.dims;
46    let n = nx * ny * nz;
47    assert_eq!(field_ppm.len(), n, "field length must match grid");
48    assert_eq!(mask.len(), n, "mask length must match grid");
49
50    // Column-major (crate) -> row-major (X,Y,Z) working volume.
51    let mut vol = vec![0.0f64; n];
52    for z in 0..nz {
53        for y in 0..ny {
54            for x in 0..nx {
55                vol[ri(x, y, z, ny, nz)] = field_ppm[x + nx * (y + ny * z)];
56            }
57        }
58    }
59
60    // Intermediate padded size so (dim - OUT) is a multiple of SHIFT.
61    let pad_to = |d: usize| -> usize {
62        let e = d - OUT;
63        e.div_ceil(SHIFT) * SHIFT - e // extra samples appended
64    };
65    let (padx, pady, padz) = (pad_to(nx), pad_to(ny), pad_to(nz));
66    let (xp, yp, zp) = (nx + padx, ny + pady, nz + padz); // output-grid size
67
68    // Edge-pad: MARGIN before, (pad + MARGIN) after, per axis (replicate).
69    let (bx, by, bz) = (MARGIN, MARGIN, MARGIN);
70    let (dx, dy, dz) = (xp + 2 * MARGIN, yp + 2 * MARGIN, zp + 2 * MARGIN);
71    let clamp = |v: isize, lo: usize, hi: usize| (v.clamp(lo as isize, hi as isize - 1)) as usize;
72    let mut padded = vec![0.0f64; dx * dy * dz];
73    for x in 0..dx {
74        let sx = clamp(x as isize - bx as isize, 0, nx);
75        for y in 0..dy {
76            let sy = clamp(y as isize - by as isize, 0, ny);
77            for z in 0..dz {
78                let sz = clamp(z as isize - bz as isize, 0, nz);
79                padded[ri(x, y, z, dy, dz)] = vol[ri(sx, sy, sz, ny, nz)];
80            }
81        }
82    }
83
84    let num_i = (xp - OUT) / SHIFT + 1;
85    let num_j = (yp - OUT) / SHIFT + 1;
86    let num_k = (zp - OUT) / SHIFT + 1;
87
88    let model = OnnxModel::load(model_onnx)?;
89    let mut output = vec![0.0f64; xp * yp * zp];
90    let mut patches: Vec<Vec<f64>> = Vec::with_capacity(num_i * num_j * num_k);
91
92    for k in 0..num_k {
93        for j in 0..num_j {
94            for i in 0..num_i {
95                // Extract the 64³ input patch (already row-major = NCDHW buffer).
96                let mut buf = vec![0.0f32; IN * IN * IN];
97                for px in 0..IN {
98                    for py in 0..IN {
99                        for pz in 0..IN {
100                            let src = ri(SHIFT * i + px, SHIFT * j + py, SHIFT * k + pz, dy, dz);
101                            buf[ri(px, py, pz, IN, IN)] = padded[src] as f32;
102                        }
103                    }
104                }
105                let out = model.run_single(&Tensor::new(vec![1, 1, IN, IN, IN], buf))?;
106                let mut patch: Vec<f64> = out.data.iter().map(|&v| v as f64).collect();
107
108                // Linear-ramp blend against already-placed neighbors.
109                if i != 0 {
110                    blend(&mut patch, &patches[patches.len() - 1], 0);
111                }
112                if j != 0 {
113                    blend(&mut patch, &patches[patches.len() - num_i], 1);
114                }
115                if k != 0 {
116                    blend(&mut patch, &patches[patches.len() - num_i * num_j], 2);
117                }
118
119                for ox in 0..OUT {
120                    for oy in 0..OUT {
121                        for oz in 0..OUT {
122                            output[ri(SHIFT * i + ox, SHIFT * j + oy, SHIFT * k + oz, yp, zp)] =
123                                patch[ri(ox, oy, oz, OUT, OUT)];
124                        }
125                    }
126                }
127                patches.push(patch);
128            }
129        }
130    }
131
132    // Crop to the input grid, mask, and repack row-major -> column-major.
133    let mut chi = vec![0.0f64; n];
134    for z in 0..nz {
135        for y in 0..ny {
136            for x in 0..nx {
137                let i = x + nx * (y + ny * z);
138                if mask[i] != 0 {
139                    chi[i] = output[ri(x, y, z, yp, zp)];
140                }
141            }
142        }
143    }
144    Ok(chi)
145}
146
147/// Blend the leading `OVERLAP` slabs of `patch` with the trailing slabs of a
148/// previously-placed `neighbor` along `dir` (0=x,1=y,2=z), ramping neighbor→patch
149/// (matches `util.patch_process`).
150fn blend(patch: &mut [f64], neighbor: &[f64], dir: usize) {
151    let denom = (OVERLAP - 1) as f64;
152    for t in 0..OVERLAP {
153        let w = 1.0 - t as f64 / denom; // weight on the neighbor
154        for a in 0..OUT {
155            for b in 0..OUT {
156                let (pi, ni) = match dir {
157                    0 => (ri(t, a, b, OUT, OUT), ri(OUT - OVERLAP + t, a, b, OUT, OUT)),
158                    1 => (ri(a, t, b, OUT, OUT), ri(a, OUT - OVERLAP + t, b, OUT, OUT)),
159                    _ => (ri(a, b, t, OUT, OUT), ri(a, b, OUT - OVERLAP + t, OUT, OUT)),
160                };
161                patch[pi] = w * neighbor[ni] + (1.0 - w) * patch[pi];
162            }
163        }
164    }
165}
qsm_core/inversion/autoqsm.rs

qsm_core/inversion/
autoqsm.rs
