qsm_core/bgremove/
bfrnet.rs

1//! BFRnet deep-learning background field removal (`onnx` feature).
2//!
3//! BFRnet is a 3D dual-frequency octave-convolution U-Net that predicts the
4//! **background** field from a masked total field; the local tissue field is
5//! `total − background`, re-masked. Everything is in ppm — the network never
6//! sees TE/B0/B0-direction. This mirrors the QSM-CI ONNX port (faithful to the
7//! authors' MATLAB `predict` to |Δ| ≈ 1e-7).
8//!
9//! The whole volume runs in one pass, as the authors run it. Its activations scale with the
10//! volume: a 184×256×128 scan needs more than a 32-bit (WASM) heap holds. [`bfrnet_tiled`] runs
11//! it patch by patch instead, with memory bounded by the patch; see it for what that costs.
12//!
13//! Weights are not bundled; the caller passes the exported `bfrnet.onnx` bytes
14//! (see [`crate::models`]).
15//!
16//! # Reference pipeline
17//!
18//! From Zhu et al., Z Med Phys 2023 (doi:10.1016/j.zemedi.2022.08.001), Methods:
19//!
20//! - **Training data:** 1 mm isotropic 3T (96 subjects, 64³ patches) with simulated healthy and
21//!   hemorrhage sources.
22//! - **In-vivo test inputs:** best-path unwrapping (Laplacian unwrapping for the hemorrhage
23//!   patient), FSL `bet2 -f 0.2` (`-w 2`) mask eroded by 3 voxels, total field from a fit of the
24//!   unwrapped phase against TE; iLSQR after BFRnet.
25//! - **Resolution:** "we have only evaluated testing datasets with 1 mm³ isotropic image resolution
26//!   since this matches the training datasets. BFRnet's application on total field maps of other
27//!   image resolutions may need retraining" (Discussion). The authors' `BFRnet_demo.m` does not
28//!   resample.
29//!
30//! # Agreement with the original
31//!
32//! `tests/models_onnx/field.rs` runs it on a 36×32×28 box of in-vivo data (1 mm, 2.89 T, across the
33//! brain's edge) against its original's verified QSM-CI entry, BFRnet (`bfrnet`) (QSMxT/QSM-CI#316;
34//! the fixture is regenerated in Actions by `tests/data/dl_field_qsmci_small.py`): NRMSE 2.1e-4 %,
35//! max |Δ| 1.6e-7 ppm (measured 2026-10-30): float32 rounding. The box is one of the original
36//! wrapper's 128³ patches, so it runs whole there, as [`bfrnet`] does here; on volumes larger than
37//! 128³ the wrapper tiles and this does not (0.56 % apart on QSM-CI's 128×160×96 case).
38//!
39//! # Pipeline rules
40//!
41//! The facts above that a configuration can depart from are published as rules in
42//! [`crate::pipeline::rules::RULES`] (each with its strength, basis, reason and source): `bfrnet.mask-bet`.
43
44use crate::grid::Grid;
45use crate::inversion::TileConfig;
46use crate::models::onnx::{OnnxModel, OnnxError, Tensor};
47
48/// Run BFRnet on a total field map.
49///
50/// * `field_ppm` — total field (ppm), column-major `(nx,ny,nz)`.
51/// * `mask` — binary brain mask (same layout).
52/// * `model_onnx` — bytes of the exported `bfrnet.onnx`.
53///
54/// Returns the local tissue field (ppm), masked, in the same layout.
55///
56/// The whole volume is run in one pass. BFRnet is fully convolutional with three
57/// pooling levels, so each spatial dimension is zero-padded up to a multiple of
58/// 8 and cropped back. For bounded memory see [`bfrnet_tiled`].
59pub fn bfrnet(
60    field_ppm: &[f64],
61    mask: &[u8],
62    grid: &Grid,
63    model_onnx: &[u8],
64) -> Result<Vec<f64>, OnnxError> {
65    let (nx, ny, nz) = grid.dims;
66    let n = nx * ny * nz;
67    assert_eq!(field_ppm.len(), n, "field length must match grid");
68    assert_eq!(mask.len(), n, "mask length must match grid");
69
70    // The net was trained on masked total field.
71    let masked: Vec<f64> = field_ppm
72        .iter()
73        .zip(mask)
74        .map(|(&f, &m)| if m != 0 { f } else { 0.0 })
75        .collect();
76
77    // Padded spatial dims (multiple of 8).
78    let (px, py, pz) = (
79        nx.div_ceil(8) * 8,
80        ny.div_ceil(8) * 8,
81        nz.div_ceil(8) * 8,
82    );
83
84    // Repack column-major (nx,ny,nz) f64 → row-major (px,py,pz) f32, zero-padded.
85    // ONNX/tract tensors are row-major (last axis fastest), so the axis order
86    // fed to the net matches the Python port's (X,Y,Z).
87    let mut input = vec![0.0f32; px * py * pz];
88    for z in 0..nz {
89        for y in 0..ny {
90            for x in 0..nx {
91                let src = x + nx * (y + ny * z);
92                let dst = z + pz * (y + py * x);
93                input[dst] = masked[src] as f32;
94            }
95        }
96    }
97
98    let model = OnnxModel::load(model_onnx)?;
99    let out = model.run_single(&Tensor::new(vec![1, 1, px, py, pz], input))?;
100    if out.shape != [1, 1, px, py, pz] {
101        return Err(OnnxError::Run(format!(
102            "unexpected output shape {:?}, expected [1,1,{px},{py},{pz}]",
103            out.shape
104        )));
105    }
106
107    // Crop back to (nx,ny,nz), form local = masked − background, re-mask,
108    // and unpack row-major → column-major.
109    let mut local = vec![0.0f64; n];
110    for z in 0..nz {
111        for y in 0..ny {
112            for x in 0..nx {
113                let dst = x + nx * (y + ny * z);
114                if mask[dst] != 0 {
115                    let src = z + pz * (y + py * x);
116                    local[dst] = masked[dst] - out.data[src] as f64;
117                }
118            }
119        }
120    }
121    Ok(local)
122}
123
124/// BFRnet patch by patch: each `core` block of the background field is predicted from a patch
125/// with `halo` voxels of context on every side (zero outside the volume), padded to a multiple of
126/// 8, so peak memory follows the patch, not the volume.
127///
128/// An approximation of [`bfrnet`]: each patch sees only its own part of the masked field, while
129/// the background field it predicts comes from sources outside the brain. Tiles that do not touch
130/// `mask` are skipped. Returns the local field (ppm), masked, like [`bfrnet`].
131pub fn bfrnet_tiled(
132    field_ppm: &[f64],
133    mask: &[u8],
134    grid: &Grid,
135    model_onnx: &[u8],
136    cfg: &TileConfig,
137    progress: impl FnMut(usize, usize),
138) -> Result<Vec<f64>, OnnxError> {
139    let n = grid.n_total();
140    assert_eq!(field_ppm.len(), n, "field length must match grid");
141    assert_eq!(mask.len(), n, "mask length must match grid");
142    let masked: Vec<f64> = field_ppm.iter().zip(mask).map(|(&f, &m)| if m != 0 { f } else { 0.0 }).collect();
143    let model = OnnxModel::load(model_onnx)?;
144    let background = crate::inversion::tiled::tiled_field_inversion(
145        &masked, mask, grid, &model, &crate::inversion::tiled::tile_cost_of("bfrnet"), cfg, |v| v as f32, |o| o as f64, progress,
146    )?;
147    Ok(masked
148        .iter()
149        .zip(&background)
150        .zip(mask)
151        .map(|((&f, &b), &m)| if m != 0 { f - b } else { 0.0 })
152        .collect())
153}
qsm_core/bgremove/bfrnet.rs

qsm_core/bgremove/
bfrnet.rs
