qsm_core/pipeline/
bg_removal.rs

1//! Background field removal stage
2//!
3//! Dispatcher that calls the appropriate background removal algorithm
4//! based on configuration, including algorithm-specific parameter defaults
5//! (V-SHARP radii computation, LBV max_iter, iSMV radius).
6
7use super::config::*;
8
9/// Run background field removal on a total field map.
10///
11/// # Arguments
12/// * `field_ppm` - Total field map in ppm
13/// * `mask` - Binary brain mask
14/// * `metadata` - Scan metadata
15/// * `config` - Background removal configuration
16/// * `progress` - Progress callback (current_iter, max_iter)
17///
18/// # Returns
19/// `BgRemovalResult` with local field in ppm and eroded mask
20pub fn run_bg_removal(
21    field_ppm: &[f64],
22    mask: &[u8],
23    metadata: &ScanMetadata,
24    config: &BgRemovalConfig,
25    progress: &mut dyn FnMut(usize, usize),
26) -> Result<BgRemovalResult, PipelineError> {
27    let grid = metadata.grid();
28    let n_voxels = grid.n_total();
29
30    if field_ppm.len() != n_voxels {
31        return Err(PipelineError::DimensionMismatch {
32            expected: n_voxels,
33            got: field_ppm.len(),
34        });
35    }
36
37    // mSMV needs the scan's field strength and echo time for its ppm↔radian cap.
38    let msmv_scan = |base: crate::bgremove::MsmvParams| crate::bgremove::MsmvParams {
39        b0: metadata.field_strength,
40        te: metadata.echo_times.first().copied().unwrap_or(base.te),
41        ..base
42    };
43
44    let (mut local_field, eroded_mask) = match config.algorithm {
45        BgRemovalAlgorithm::Vsharp => {
46            crate::bgremove::vsharp(
47                field_ppm, mask, &grid, &config.vsharp, progress,
48            )
49        }
50        BgRemovalAlgorithm::Pdf => {
51            let local = crate::bgremove::pdf(
52                field_ppm, mask, &grid,
53                metadata.b0_direction, &config.pdf, progress,
54            );
55            (local, mask.to_vec())
56        }
57        BgRemovalAlgorithm::Lbv => {
58            crate::bgremove::lbv(
59                field_ppm, mask, &grid, &config.lbv, progress,
60            )
61        }
62        BgRemovalAlgorithm::Ismv => {
63            crate::bgremove::ismv(
64                field_ppm, mask, &grid, &config.ismv, progress,
65            )
66        }
67        BgRemovalAlgorithm::Sharp => {
68            crate::bgremove::sharp(
69                field_ppm, mask, &grid, &config.sharp,
70            )
71        }
72        BgRemovalAlgorithm::Resharp => {
73            crate::bgremove::resharp(
74                field_ppm, mask, &grid,
75                &config.resharp,
76                progress,
77            )
78        }
79        BgRemovalAlgorithm::Harperella => {
80            crate::bgremove::harperella(
81                field_ppm, mask, &grid,
82                &config.harperella,
83                progress,
84            )
85        }
86        BgRemovalAlgorithm::Iharperella => {
87            crate::bgremove::iharperella(
88                field_ppm, mask, &grid,
89                &config.harperella,
90                progress,
91            )
92        }
93        BgRemovalAlgorithm::Bfrnet => {
94            let local = run_bfrnet(field_ppm, mask, &grid)?;
95            (local, mask.to_vec())
96        }
97    };
98
99    // Optional mSMV boundary-shadow refinement of the primary BFR's local field.
100    // mSMV is a refinement (not a standalone primary remover), so this post-step
101    // is the only way it's wired into the pipeline (Roberts 2024).
102    if config.msmv_refine {
103        let params = msmv_scan(crate::bgremove::MsmvParams {
104            prefilter: false,
105            ..config.msmv.clone()
106        });
107        let (refined, _) = crate::bgremove::msmv(&local_field, &eroded_mask, &grid, &params, |_, _| {});
108        local_field = refined;
109    }
110
111    Ok(BgRemovalResult {
112        local_field_ppm: local_field,
113        eroded_mask,
114    })
115}
116
117/// Source the BFRnet weights and run inference. Requires the `onnx` feature;
118/// weights come from the model registry (local `$QSM_MODEL_DIR`/cache, or the
119/// `download` feature).
120#[cfg(feature = "onnx")]
121fn run_bfrnet(
122    field_ppm: &[f64],
123    mask: &[u8],
124    grid: &crate::Grid,
125) -> Result<Vec<f64>, PipelineError> {
126    let spec = crate::models::find_model("bfrnet")
127        .ok_or_else(|| PipelineError::InvalidConfig("bfrnet not in model registry".into()))?;
128    let bytes = crate::models::primary_weight_bytes(spec)
129        .map_err(PipelineError::InvalidConfig)?;
130    crate::bgremove::bfrnet(field_ppm, mask, grid, &bytes)
131        .map_err(|e| PipelineError::AlgorithmError(e.to_string()))
132}
133
134#[cfg(not(feature = "onnx"))]
135fn run_bfrnet(
136    _field_ppm: &[f64],
137    _mask: &[u8],
138    _grid: &crate::Grid,
139) -> Result<Vec<f64>, PipelineError> {
140    Err(PipelineError::InvalidConfig(
141        "BFRnet requires building qsm-core with the 'onnx' feature".into(),
142    ))
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148
149    #[test]
150    fn test_bg_removal_dispatches_vsharp() {
151        let (nx, ny, nz) = (8, 8, 8);
152        let n = nx * ny * nz;
153        let field = vec![0.1; n];
154        let mask = vec![1u8; n];
155        let meta = ScanMetadata {
156            dims: (nx, ny, nz),
157            voxel_size: (1.0, 1.0, 1.0),
158            echo_times: vec![0.005],
159            field_strength: 3.0,
160            b0_direction: (0.0, 0.0, 1.0),
161        };
162
163        let result = run_bg_removal(
164            &field, &mask, &meta, &BgRemovalConfig::default(),
165            &mut |_, _| {},
166        );
167        assert!(result.is_ok());
168        let r = result.unwrap();
169        assert_eq!(r.local_field_ppm.len(), n);
170        assert_eq!(r.eroded_mask.len(), n);
171    }
172
173    #[test]
174    fn test_bg_removal_dispatches_pdf() {
175        let (nx, ny, nz) = (8, 8, 8);
176        let n = nx * ny * nz;
177        let field = vec![0.1; n];
178        let mask = vec![1u8; n];
179        let meta = ScanMetadata {
180            dims: (nx, ny, nz), voxel_size: (1.0, 1.0, 1.0),
181            echo_times: vec![0.005], field_strength: 3.0, b0_direction: (0.0, 0.0, 1.0),
182        };
183        let config = BgRemovalConfig { algorithm: BgRemovalAlgorithm::Pdf, ..Default::default() };
184        let r = run_bg_removal(&field, &mask, &meta, &config, &mut |_, _| {}).unwrap();
185        assert_eq!(r.local_field_ppm.len(), n);
186    }
187
188    /// PDF is the one background remover that builds a dipole kernel, so the
189    /// pipeline must hand it `metadata.b0_direction` rather than assuming `+z`.
190    /// Hardcoding `(0.0, 0.0, 1.0)` makes both calls below identical.
191    #[test]
192    fn test_bg_removal_pdf_uses_metadata_b0_direction() {
193        let (nx, ny, nz) = (12, 12, 12);
194        let n = nx * ny * nz;
195
196        // Interior mask with real background voxels around it, so PDF has
197        // background sources to project onto, plus a spatially varying field.
198        let mut field = vec![0.0f64; n];
199        let mut mask = vec![0u8; n];
200        for k in 0..nz {
201            for j in 0..ny {
202                for i in 0..nx {
203                    let idx = i + j * nx + k * nx * ny;
204                    let (x, y, z) = (i as f64 - 5.5, j as f64 - 5.5, k as f64 - 5.5);
205                    if x * x + y * y + z * z < 9.0 {
206                        mask[idx] = 1;
207                    }
208                    field[idx] = 0.01 * (0.3 * x + 0.5 * y - 0.7 * z + 0.02 * x * y * z);
209                }
210            }
211        }
212
213        let config = BgRemovalConfig { algorithm: BgRemovalAlgorithm::Pdf, ..Default::default() };
214        let run = |bdir: (f64, f64, f64)| {
215            let meta = ScanMetadata {
216                dims: (nx, ny, nz), voxel_size: (1.0, 1.0, 1.0),
217                echo_times: vec![0.005], field_strength: 3.0, b0_direction: bdir,
218            };
219            run_bg_removal(&field, &mask, &meta, &config, &mut |_, _| {}).unwrap().local_field_ppm
220        };
221
222        let axial = run((0.0, 0.0, 1.0));
223        // ~30 degrees off +z in the y-z plane.
224        let theta = std::f64::consts::FRAC_PI_6;
225        let oblique = run((0.0, theta.sin(), theta.cos()));
226
227        assert_eq!(axial.len(), n);
228        assert_eq!(oblique.len(), n);
229
230        let max_diff = axial
231            .iter()
232            .zip(&oblique)
233            .map(|(a, b)| (a - b).abs())
234            .fold(0.0f64, f64::max);
235        assert!(
236            max_diff > 1e-9,
237            "oblique B0 gave the same local field as axial (max diff {max_diff:e}); \
238             b0_direction is not reaching the PDF dipole kernel",
239        );
240    }
241
242    #[test]
243    fn test_bg_removal_dispatches_lbv() {
244        let (nx, ny, nz) = (8, 8, 8);
245        let n = nx * ny * nz;
246        let field = vec![0.1; n];
247        let mask = vec![1u8; n];
248        let meta = ScanMetadata {
249            dims: (nx, ny, nz), voxel_size: (1.0, 1.0, 1.0),
250            echo_times: vec![0.005], field_strength: 3.0, b0_direction: (0.0, 0.0, 1.0),
251        };
252        let config = BgRemovalConfig { algorithm: BgRemovalAlgorithm::Lbv, ..Default::default() };
253        let r = run_bg_removal(&field, &mask, &meta, &config, &mut |_, _| {}).unwrap();
254        assert_eq!(r.local_field_ppm.len(), n);
255    }
256
257    #[test]
258    fn test_bg_removal_dispatches_sharp() {
259        let (nx, ny, nz) = (8, 8, 8);
260        let n = nx * ny * nz;
261        let field = vec![0.1; n];
262        let mask = vec![1u8; n];
263        let meta = ScanMetadata {
264            dims: (nx, ny, nz), voxel_size: (1.0, 1.0, 1.0),
265            echo_times: vec![0.005], field_strength: 3.0, b0_direction: (0.0, 0.0, 1.0),
266        };
267        let config = BgRemovalConfig { algorithm: BgRemovalAlgorithm::Sharp, ..Default::default() };
268        let r = run_bg_removal(&field, &mask, &meta, &config, &mut |_, _| {}).unwrap();
269        assert_eq!(r.local_field_ppm.len(), n);
270    }
271
272    #[test]
273    fn test_bg_removal_dispatches_ismv() {
274        let (nx, ny, nz) = (8, 8, 8);
275        let n = nx * ny * nz;
276        let field = vec![0.1; n];
277        let mask = vec![1u8; n];
278        let meta = ScanMetadata {
279            dims: (nx, ny, nz), voxel_size: (1.0, 1.0, 1.0),
280            echo_times: vec![0.005], field_strength: 3.0, b0_direction: (0.0, 0.0, 1.0),
281        };
282        let config = BgRemovalConfig { algorithm: BgRemovalAlgorithm::Ismv, ..Default::default() };
283        let r = run_bg_removal(&field, &mask, &meta, &config, &mut |_, _| {}).unwrap();
284        assert_eq!(r.local_field_ppm.len(), n);
285    }
286
287    #[test]
288    fn test_bg_removal_msmv_refine_post_step() {
289        let (nx, ny, nz) = (8, 8, 8);
290        let n = nx * ny * nz;
291        let field = vec![0.1; n];
292        let mask = vec![1u8; n];
293        let meta = ScanMetadata {
294            dims: (nx, ny, nz), voxel_size: (1.0, 1.0, 1.0),
295            echo_times: vec![0.005], field_strength: 3.0, b0_direction: (0.0, 0.0, 1.0),
296        };
297        // Primary V-SHARP + mSMV boundary-shadow refinement.
298        let config = BgRemovalConfig {
299            algorithm: BgRemovalAlgorithm::Vsharp,
300            msmv_refine: true,
301            ..Default::default()
302        };
303        let r = run_bg_removal(&field, &mask, &meta, &config, &mut |_, _| {}).unwrap();
304        assert_eq!(r.local_field_ppm.len(), n);
305    }
306
307    #[test]
308    fn test_bg_removal_dispatches_resharp() {
309        let (nx, ny, nz) = (8, 8, 8);
310        let n = nx * ny * nz;
311        let field = vec![0.1; n];
312        let mask = vec![1u8; n];
313        let meta = ScanMetadata {
314            dims: (nx, ny, nz), voxel_size: (1.0, 1.0, 1.0),
315            echo_times: vec![0.005], field_strength: 3.0, b0_direction: (0.0, 0.0, 1.0),
316        };
317        let config = BgRemovalConfig { algorithm: BgRemovalAlgorithm::Resharp, ..Default::default() };
318        let r = run_bg_removal(&field, &mask, &meta, &config, &mut |_, _| {}).unwrap();
319        assert_eq!(r.local_field_ppm.len(), n);
320    }
321
322    #[test]
323    fn test_bg_removal_validates_dims() {
324        let meta = ScanMetadata {
325            dims: (4, 4, 4),
326            voxel_size: (1.0, 1.0, 1.0),
327            echo_times: vec![0.005],
328            field_strength: 3.0,
329            b0_direction: (0.0, 0.0, 1.0),
330        };
331        let field = vec![0.0; 32]; // wrong size, should be 64
332        let mask = vec![1u8; 64];
333
334        let result = run_bg_removal(
335            &field, &mask, &meta, &BgRemovalConfig::default(),
336            &mut |_, _| {},
337        );
338        assert!(result.is_err());
339    }
340}
qsm_core/pipeline/bg_removal.rs

qsm_core/pipeline/
bg_removal.rs
