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//! Weights are not bundled; the caller passes the exported `bfrnet.onnx` bytes
10//! (see [`crate::models`]).
11
12use crate::grid::Grid;
13use crate::models::onnx::{OnnxModel, OnnxError, Tensor};
14
15/// Run BFRnet on a total field map.
16///
17/// * `field_ppm` — total field (ppm), column-major `(nx,ny,nz)`.
18/// * `mask` — binary brain mask (same layout).
19/// * `model_onnx` — bytes of the exported `bfrnet.onnx`.
20///
21/// Returns the local tissue field (ppm), masked, in the same layout.
22///
23/// The whole volume is run in one pass. BFRnet is fully convolutional with three
24/// pooling levels, so each spatial dimension is zero-padded up to a multiple of
25/// 8 and cropped back. (Memory-bounded Hann-blended tiling — as the Python port
26/// uses for very large volumes — is a future refinement.)
27pub fn bfrnet(
28    field_ppm: &[f64],
29    mask: &[u8],
30    grid: &Grid,
31    model_onnx: &[u8],
32) -> Result<Vec<f64>, OnnxError> {
33    let (nx, ny, nz) = grid.dims;
34    let n = nx * ny * nz;
35    assert_eq!(field_ppm.len(), n, "field length must match grid");
36    assert_eq!(mask.len(), n, "mask length must match grid");
37
38    // The net was trained on masked total field.
39    let masked: Vec<f64> = field_ppm
40        .iter()
41        .zip(mask)
42        .map(|(&f, &m)| if m != 0 { f } else { 0.0 })
43        .collect();
44
45    // Padded spatial dims (multiple of 8).
46    let (px, py, pz) = (
47        nx.div_ceil(8) * 8,
48        ny.div_ceil(8) * 8,
49        nz.div_ceil(8) * 8,
50    );
51
52    // Repack column-major (nx,ny,nz) f64 → row-major (px,py,pz) f32, zero-padded.
53    // ONNX/tract tensors are row-major (last axis fastest), so the axis order
54    // fed to the net matches the Python port's (X,Y,Z).
55    let mut input = vec![0.0f32; px * py * pz];
56    for z in 0..nz {
57        for y in 0..ny {
58            for x in 0..nx {
59                let src = x + nx * (y + ny * z);
60                let dst = z + pz * (y + py * x);
61                input[dst] = masked[src] as f32;
62            }
63        }
64    }
65
66    let model = OnnxModel::load(model_onnx)?;
67    let out = model.run_single(&Tensor::new(vec![1, 1, px, py, pz], input))?;
68    if out.shape != [1, 1, px, py, pz] {
69        return Err(OnnxError::Run(format!(
70            "unexpected output shape {:?}, expected [1,1,{px},{py},{pz}]",
71            out.shape
72        )));
73    }
74
75    // Crop back to (nx,ny,nz), form local = masked − background, re-mask,
76    // and unpack row-major → column-major.
77    let mut local = vec![0.0f64; n];
78    for z in 0..nz {
79        for y in 0..ny {
80            for x in 0..nx {
81                let dst = x + nx * (y + ny * z);
82                if mask[dst] != 0 {
83                    let src = z + pz * (y + py * x);
84                    local[dst] = masked[dst] - out.data[src] as f64;
85                }
86            }
87        }
88    }
89    Ok(local)
90}
qsm_core/bgremove/bfrnet.rs

qsm_core/bgremove/
bfrnet.rs
