qsm_core/models/
registry.rs

1//! The static table of known deep-learning QSM models.
2//!
3//! Weights are hosted externally and fetched on use (see [`super`]); nothing is
4//! vendored here. Entries marked [`WeightStatus::Pending`] are recognized
5//! targets whose ONNX weights are not yet converted and hosted — their `url`
6//! and `sha256` are filled in once the exported `.onnx` is uploaded (Hugging Face). Until then a host can still run them via bring-your-own-weights
7//! (`$QSM_MODEL_DIR`).
8//!
9//! Editing checklist when a model goes from Pending → Available:
10//! 1. export/convert to ONNX, upload the file,
11//! 2. set `url`, `sha256` (lowercase hex), `bytes`,
12//! 3. flip `status` to [`WeightStatus::Available`],
13//! 4. confirm `inputs`/`outputs`/`size_divisor` match the exported graph.
14
15use super::{Framework, ModelSpec, ModelStage, WeightFile, WeightStatus};
16
17/// All models known to QSM-Core, in a stable order.
18pub fn all_models() -> &'static [ModelSpec] {
19    MODELS
20}
21
22/// Look up a model by its [`ModelSpec::id`] (case-sensitive), e.g. `"qsmnet"`.
23pub fn find_model(id: &str) -> Option<&'static ModelSpec> {
24    MODELS.iter().find(|m| m.id == id)
25}
26
27// A single ONNX weight file whose location is not yet known (Pending model).
28const fn pending_onnx(name: &'static str) -> WeightFile {
29    WeightFile { name, url: "", sha256: "", bytes: 0 }
30}
31
32const MODELS: &[ModelSpec] = &[
33    // ---- Background field removal ------------------------------------------
34    ModelSpec {
35        id: "bfrnet",
36        name: "BFRnet",
37        stage: ModelStage::BackgroundRemoval,
38        status: WeightStatus::Available,
39        origin: Framework::Matlab,
40        description: "Dual-frequency octave-convolution U-Net for background field \
41                      removal (total field → local field). Fully convolutional.",
42        paper: "Kames et al. / Sun group; https://github.com/sunhongfu/BFRnet",
43        source: "https://github.com/sunhongfu/BFRnet",
44        license: "",
45        // Mirrored on Hugging Face (qsmxt/qsm-onnx-weights). Verified: anonymous
46        // download + SHA-256 match.
47        files: &[WeightFile {
48            name: "bfrnet.onnx",
49            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/bfrnet.onnx",
50            sha256: "6f693f0a02c94550179c4b5188ce652fc4bd8198ff57aaddd102fba67fe873d7",
51            bytes: 79_600_612,
52        }],
53        inputs: &["field"],
54        outputs: &["local_field"],
55        size_divisor: 8,
56    },
57    // ---- Dipole inversion --------------------------------------------------
58    ModelSpec {
59        id: "xqsm",
60        name: "xQSM",
61        stage: ModelStage::DipoleInversion,
62        status: WeightStatus::Available,
63        origin: Framework::PyTorch,
64        description: "Octave-convolution U-Net with a learned residual for dipole \
65                      inversion (local field → susceptibility).",
66        paper: "Gao et al., NMR Biomed 2021; doi:10.1002/nbm.4461",
67        source: "https://github.com/sunhongfu/xQSM",
68        license: "",
69        // Exported from xQSM_invivo.pth (v1.0-demo) with scripts/onnx-export/export_xqsm.py;
70        // torch↔onnxruntime parity max|Δ| ≈ 6e-5. Mirrored on Hugging Face (qsmxt/qsm-onnx-weights).
71        files: &[WeightFile {
72            name: "xqsm.onnx",
73            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/xqsm.onnx",
74            sha256: "81854ec2ca85bba25c2f9aae05efdcac299d0efffb03fe6a0e6797019e46b487",
75            bytes: 20_901_826,
76        }],
77        inputs: &["field"],
78        outputs: &["chi"],
79        size_divisor: 8,
80    },
81    ModelSpec {
82        id: "qsmnet",
83        name: "QSMnet",
84        stage: ModelStage::DipoleInversion,
85        status: WeightStatus::Available,
86        origin: Framework::TensorFlow,
87        description: "3D U-Net dipole inversion trained on COSMOS. Expects 1 mm \
88                      isotropic input; z-scored with training mean/std.",
89        paper: "Yoon et al., NeuroImage 2018; doi:10.1016/j.neuroimage.2018.06.030",
90        source: "https://github.com/SNU-LIST/QSMnet",
91        license: "",
92        // Clean PyTorch re-export of the TF1.14 checkpoint (tract-friendly NCDHW);
93        // see scripts/onnx-export/export_qsmnet.py. Mirrored on Hugging Face (qsmxt/qsm-onnx-weights).
94        files: &[WeightFile {
95            name: "qsmnet.onnx",
96            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/qsmnet.onnx",
97            sha256: "8fd1d79b7a9258a262a2faab9e8469757ae7735f3da6212e431acfb207c21b54",
98            bytes: 397_801_854,
99        }],
100        inputs: &["field"],
101        outputs: &["chi"],
102        size_divisor: 16,
103    },
104    ModelSpec {
105        id: "qsmnet-plus",
106        name: "QSMnet+",
107        stage: ModelStage::DipoleInversion,
108        status: WeightStatus::Available,
109        origin: Framework::TensorFlow,
110        description: "QSMnet retrained with susceptibility-scaling augmentation for \
111                      a wider, more linear χ range.",
112        paper: "Jung et al., NeuroImage 2020; doi:10.1016/j.neuroimage.2020.116579",
113        source: "https://github.com/SNU-LIST/QSMnet",
114        license: "",
115        // Clean PyTorch re-export of the TF1.14 QSMnet+_64 checkpoint (same U-Net
116        // as QSMnet, different weights + norm). Mirrored on Hugging Face (qsmxt/qsm-onnx-weights).
117        files: &[WeightFile {
118            name: "qsmnet-plus.onnx",
119            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/qsmnet-plus.onnx",
120            sha256: "0ed31f9a1b66f75fee4a96022b6bc3838b91bc5b819714dade5ca87cd331683d",
121            bytes: 397_801_854,
122        }],
123        inputs: &["field"],
124        outputs: &["chi"],
125        size_divisor: 16,
126    },
127    ModelSpec {
128        id: "qsmgan",
129        name: "QSMGAN",
130        stage: ModelStage::DipoleInversion,
131        status: WeightStatus::Available,
132        origin: Framework::PyTorch,
133        description: "3D U-Net generator (WGAN-GP refined) for dipole inversion, run \
134                      patch-wise (64³ input → 48³ output, i64o48). Only the generator is \
135                      used at inference; the sign flip, input_scale/tanh (χ=atanh/10) and \
136                      patch tiling live in the Rust glue.",
137        paper: "Chen et al., NeuroImage 2020; doi:10.1016/j.neuroimage.2019.116389",
138        source: "https://github.com/mmorri10/QSMGAN-LupoLab",
139        license: "MIT (fork)",
140        files: &[WeightFile {
141            name: "qsmgan.onnx",
142            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/qsmgan.onnx",
143            sha256: "238573a6fa8a563004fc28e63a30694289d1a1a204f3c406daba486b1122d604",
144            bytes: 10_753_955,
145        }],
146        inputs: &["field"],
147        outputs: &["chi"],
148        size_divisor: 8,
149    },
150    ModelSpec {
151        id: "lpcnn",
152        name: "LPCNN",
153        stage: ModelStage::DipoleInversion,
154        status: WeightStatus::Available,
155        origin: Framework::PyTorch,
156        description: "Learned proximal CNN, 3 unrolled iterations of proximal \
157                      gradient descent. The k-space dipole data-consistency step, the \
158                      unroll, the learned step size and the mean/std normalization run \
159                      in Rust (rustfft); only the learned proximal CNN is ONNX.",
160        paper: "Lai et al., MICCAI 2020; doi:10.1007/978-3-030-59713-9_13",
161        source: "https://github.com/Sulam-Group/LPCNN",
162        license: "",
163        files: &[WeightFile {
164            name: "lpcnn.onnx",
165            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/lpcnn.onnx",
166            sha256: "831f2a161f6ab8ee7bf64cc2dd8bca532fc971d0083cf5729d6302d3ed31625d",
167            bytes: 1_790_869,
168        }],
169        inputs: &["field", "mask", "b_vec"],
170        outputs: &["chi"],
171        size_divisor: 8,
172    },
173    ModelSpec {
174        id: "ir2qsm",
175        name: "IR2QSM",
176        stage: ModelStage::DipoleInversion,
177        status: WeightStatus::Available,
178        origin: Framework::PyTorch,
179        description: "IR2U-net dipole inversion: a 3D U-net (depth 4) run for 4 \
180                      unrolled iterations with reverse concatenations and a recurrent \
181                      middle module. The whole net is ONNX; the Rust glue only does \
182                      /8 zero-pad, crop and mask (no normalization, ppm in/out). The \
183                      ungated inference-time AddNoise is pinned to its noise-free \
184                      branch for deterministic export.",
185        paper: "Li et al., Med. Phys. 2025; doi:10.1002/mp.17747; arXiv:2406.12300",
186        source: "https://github.com/YangGaoUQ/IR2QSM",
187        license: "",
188        files: &[WeightFile {
189            name: "ir2qsm.onnx",
190            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/ir2qsm.onnx",
191            sha256: "dcaf7a26d6633900f339f94be16d580ccea6147bbf6da26eacb9a5e1acdc0c53",
192            bytes: 40_259_354,
193        }],
194        inputs: &["field", "mask"],
195        outputs: &["chi"],
196        size_divisor: 8,
197    },
198    // ---- Single-step (phase/total field → χ) -------------------------------
199    ModelSpec {
200        id: "autoqsm",
201        name: "AutoQSM",
202        stage: ModelStage::SingleStep,
203        status: WeightStatus::Available,
204        origin: Framework::TensorFlow,
205        description: "V-Net single-step reconstruction (total field → susceptibility) \
206                      with no separate brain extraction. Fixed 64³→32³ patch net; \
207                      the Rust glue does the sliding-window tiling + blend.",
208        paper: "Wei et al., NeuroImage 2019; doi:10.1016/j.neuroimage.2019.116064",
209        source: "https://github.com/AMRI-Lab/AutoQSM",
210        license: "",
211        // Clean PyTorch re-export of the Keras V-Net (tract-friendly NCDHW);
212        // see scripts/onnx-export/export_autoqsm.py. Mirrored on Hugging Face (qsmxt/qsm-onnx-weights).
213        files: &[WeightFile {
214            name: "autoqsm.onnx",
215            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/autoqsm.onnx",
216            sha256: "c2f23c6ae735fd1677e8623da22cb6f25b4a932ae50ba1ef08a75439e240977e",
217            bytes: 5_609_662,
218        }],
219        // Fixed-size patch net; `size_divisor` is not a whole-volume constraint here
220        // (tiling handles padding), left at 1.
221        inputs: &["field"],
222        outputs: &["chi"],
223        size_divisor: 1,
224    },
225    ModelSpec {
226        id: "iqsm",
227        name: "iQSM",
228        stage: ModelStage::SingleStep,
229        status: WeightStatus::Available,
230        origin: Framework::PyTorch,
231        description: "LoT-Unet single-step reconstruction from wrapped phase to \
232                      susceptibility. The learnable-Laplacian front-end is fused into \
233                      the exported graph; inputs are phase, mask, TE (s), B0 (T).",
234        paper: "Gao et al., NeuroImage 2022; doi:10.1016/j.neuroimage.2022.119410",
235        source: "https://github.com/sunhongfu/iQSM",
236        license: "",
237        files: &[WeightFile {
238            name: "iqsm.onnx",
239            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/iqsm.onnx",
240            sha256: "8538a33892a877812e6cd0e22927a5040827c126e9347073d3dd0934bf226b09",
241            bytes: 17_233_763,
242        }],
243        inputs: &["phase", "mask", "te", "b0", "border"],
244        outputs: &["chi"],
245        size_divisor: 8,
246    },
247    ModelSpec {
248        id: "iqsm-plus",
249        name: "iQSM+",
250        stage: ModelStage::SingleStep,
251        status: WeightStatus::Available,
252        origin: Framework::PyTorch,
253        description: "iQSM with orientation-adaptive latent feature editing (OA-LFE); \
254                      the B0 direction is a genuine network input. Inputs: phase, \
255                      mask, TE, B0, z_prjs (B0 dir), border.",
256        paper: "Gao et al., Med Image Anal 2024; doi:10.1016/j.media.2024.103160",
257        source: "https://github.com/sunhongfu/iQSM_Plus",
258        license: "",
259        files: &[WeightFile {
260            name: "iqsm-plus.onnx",
261            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/iqsm-plus.onnx",
262            sha256: "f93a4c83804759a4a0d24603f7754efb997ddd2ba5e3a0346ed6493725cafb1d",
263            bytes: 17_740_526,
264        }],
265        inputs: &["phase", "mask", "te", "b0", "z_prjs", "border"],
266        outputs: &["chi"],
267        size_divisor: 16,
268    },
269    ModelSpec {
270        id: "iqfm",
271        name: "iQFM",
272        stage: ModelStage::PhaseToField,
273        status: WeightStatus::Available,
274        origin: Framework::PyTorch,
275        description: "The tissue-field head of the iQSM LoT-Unet: wrapped phase → \
276                      local (background-removed) field in one network (joint unwrap + \
277                      BFR). Same architecture/inputs as iQSM, `lfs` weights; output is \
278                      the local field (ppm), not susceptibility.",
279        paper: "Gao et al., NeuroImage 2022; doi:10.1016/j.neuroimage.2022.119410",
280        source: "https://github.com/sunhongfu/iQSM",
281        license: "",
282        files: &[WeightFile {
283            name: "iqfm.onnx",
284            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/iqfm.onnx",
285            sha256: "81965f4c0981612d979977c9618847e4cb1d340efa287b8dde894bdea5bf6ac9",
286            bytes: 17_233_763,
287        }],
288        inputs: &["phase", "mask", "te", "b0", "border"],
289        outputs: &["localfield"],
290        size_divisor: 8,
291    },
292    ModelSpec {
293        id: "nextqsm",
294        name: "NeXtQSM",
295        stage: ModelStage::SingleStep,
296        status: WeightStatus::Available,
297        origin: Framework::TensorFlow,
298        description: "U-Net background removal followed by a 6-step variational-network \
299                      dipole inversion (total field → susceptibility). Reimplemented as \
300                      a Rust hybrid: two exported U-Nets (`nextqsm-bf` = BFR forward, \
301                      `nextqsm-vjp` = the regularizer gradient ∇ₓ mean|VarNet(x)| written \
302                      as a forward graph) plus the FFT data-consistency gradient and unroll \
303                      in Rust. Order the files BFR-first, VJP-second.",
304        paper: "Cognolato et al., NeuroImage 2023; doi:10.1016/j.neuroimage.2022.119729",
305        source: "https://github.com/wayne1123/NeXtQSM",
306        license: "MIT",
307        files: &[
308            WeightFile {
309                name: "nextqsm-bf.onnx",
310            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/nextqsm-bf.onnx",
311                sha256: "d0f0b0391153f8e5e07ef85d46fbe492afc8394fab008350540a6868d0b4a3ba",
312                bytes: 113_211_474,
313            },
314            WeightFile {
315                name: "nextqsm-vjp.onnx",
316            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/nextqsm-vjp.onnx",
317                sha256: "50fe388663c05b1d3996139975270e500e13feb025f1b792576dbc0da814a4bc",
318                bytes: 42_437_407,
319            },
320        ],
321        inputs: &["field", "mask", "b_vec"],
322        outputs: &["chi"],
323        size_divisor: 64,
324    },
325    ModelSpec {
326        id: "modl-qsm",
327        name: "MoDL-QSM",
328        stage: ModelStage::DipoleInversion,
329        status: WeightStatus::Available,
330        origin: Framework::TensorFlow,
331        description: "Model-based deep learning: 3-iteration unroll of a QSM \
332                      forward-model gradient descent with a learned 2-channel CNN \
333                      prior. The k-space dipole data-consistency (A/Aᴴ), the unroll, \
334                      the learned step size and per-channel mean/std normalization run \
335                      in Rust (rustfft); only the CNN prior is ONNX. Output is the STI \
336                      χ33 component.",
337        paper: "Feng et al., NeuroImage 2021; doi:10.1016/j.neuroimage.2021.118376",
338        source: "https://github.com/Ruimin-Feng/MoDL-QSM",
339        license: "",
340        files: &[WeightFile {
341            name: "modl-qsm.onnx",
342            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/modl-qsm.onnx",
343            sha256: "1f8da7d39eaaa4c3f0f9e67298fba5f14230e769cfd6349c81d0b74a6f800c57",
344            bytes: 1_794_143,
345        }],
346        inputs: &["field", "mask", "b_vec"],
347        outputs: &["chi"],
348        size_divisor: 1,
349    },
350    // ---- χ-separation (χ+ / χ−) --------------------------------------------
351    ModelSpec {
352        id: "susep-net",
353        name: "SUSEP-Net",
354        stage: ModelStage::ChiSeparation,
355        status: WeightStatus::Available,
356        origin: Framework::PyTorch,
357        description: "Dual-branch 3D U-Net source separation from z-scored \
358                      [QSM, R2', local field] → [χ+, χ−]. Normalization constants \
359                      are baked into the Rust glue (SusepNetNorm).",
360        paper: "Li, Gao, Sun et al., arXiv:2506.13293 (2025)",
361        source: "https://github.com/YangGaoUQ/SUSEP-Net",
362        license: "",
363        files: &[WeightFile {
364            name: "susep-net.onnx",
365            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/susep-net.onnx",
366            sha256: "9bc85aa28451fda2c661de4b8fae8604ced7a54f028f94ed33629f6111a48b03",
367            bytes: 205_743_584,
368        }],
369        inputs: &["qsm", "r2prime", "lfs"],
370        outputs: &["chi_pos", "chi_neg"],
371        size_divisor: 8,
372    },
373    ModelSpec {
374        id: "chi-sepnet",
375        name: "χ-sepnet",
376        stage: ModelStage::ChiSeparation,
377        status: WeightStatus::Available,
378        origin: Framework::Onnx,
379        description: "SNU-LIST χ-separation network (already ONNX): a 192×192×128 \
380                      3D U-Net mapping [QSM, local field, R2′/Dr] (z-scored) → [χ+, χ−], \
381                      run as an overlapping sliding window. Normalization constants (Dr=114) \
382                      are baked into the Rust glue (ChiSepNetNorm).",
383        paper: "Kim et al. / SNU-LIST chi-separation toolbox",
384        source: "https://github.com/SNU-LIST/chi-separation",
385        license: "",
386        files: &[WeightFile {
387            name: "chi-sepnet.onnx",
388            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/chi-sepnet.onnx",
389            sha256: "4f2343649cf36b4c9b371fc732da600720060bf2e76ed845578f16fe3411b3f2",
390            bytes: 90_314_172,
391        }],
392        inputs: &["local_field", "qsm", "r2prime"],
393        outputs: &["chi_pos", "chi_neg"],
394        size_divisor: 8,
395    },
396    ModelSpec {
397        id: "hd-bet",
398        name: "HD-BET",
399        stage: ModelStage::BrainExtraction,
400        status: WeightStatus::Available,
401        origin: Framework::PyTorch,
402        description: "nnU-Net v2 3D U-Net brain extraction (HD-BET v2): magnitude → brain \
403                      mask, trained on 11,751 multi-sequence clinical MRIs. Runs at 1 mm \
404                      with Gaussian-weighted 96×192×192 sliding-window patches; nnU-Net's \
405                      crop/normalise/resample pipeline is in the Rust glue (bet::hd_bet). \
406                      Spatial axes are dynamic (patch dims multiples of 16×32×32).",
407        paper: "Isensee et al., Hum Brain Mapp 40(17):4952-4964 (2019); https://doi.org/10.1002/hbm.24750",
408        source: "https://github.com/MIC-DKFZ/HD-BET",
409        license: "CC-BY-NC-4.0",
410        files: &[WeightFile {
411            name: "hd-bet.onnx",
412            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/hd-bet.onnx",
413            sha256: "f15200f9a0697cf53b151e316a0c4fe2cacfdbe2220deb36cb56184be4a227f9",
414            bytes: 123_168_028,
415        }],
416        inputs: &["image"],
417        outputs: &["logits"],
418        size_divisor: 32,
419    },
420];
421
422#[cfg(test)]
423mod tests {
424    use super::*;
425
426    #[test]
427    fn ids_are_unique_and_findable() {
428        let mut seen = std::collections::HashSet::new();
429        for m in MODELS {
430            assert!(seen.insert(m.id), "duplicate model id: {}", m.id);
431            assert!(find_model(m.id).is_some(), "find_model failed for {}", m.id);
432            assert!(!m.name.is_empty(), "{} has empty name", m.id);
433            assert!(!m.files.is_empty(), "{} has no weight files", m.id);
434        }
435        assert!(find_model("does-not-exist").is_none());
436    }
437
438    #[test]
439    fn pending_models_are_not_available() {
440        // A Pending model is not runnable via download: it must lack a URL on at
441        // least one file (a known SHA-256 ahead of hosting is fine, e.g. BFRnet).
442        for m in MODELS.iter().filter(|m| m.status == WeightStatus::Pending) {
443            assert!(!m.is_available(), "{} is Pending but reports available", m.id);
444            assert!(
445                m.files.iter().any(|f| f.url.is_empty()),
446                "{} is Pending but every file has a url — flip status to Available",
447                m.id
448            );
449            // Any precomputed hash must still be a valid 64-char hex digest.
450            for f in m.files.iter().filter(|f| !f.sha256.is_empty()) {
451                assert_eq!(f.sha256.len(), 64, "{}/{} sha256 must be 64 hex chars", m.id, f.name);
452                assert!(
453                    f.sha256.bytes().all(|b| b.is_ascii_hexdigit()),
454                    "{}/{} sha256 not hex",
455                    m.id,
456                    f.name
457                );
458            }
459        }
460    }
461
462    #[test]
463    fn available_models_are_fully_specified() {
464        // Invariant: anything marked Available must have url + hex sha256 per file.
465        for m in MODELS.iter().filter(|m| m.status == WeightStatus::Available) {
466            for f in m.files {
467                assert!(!f.url.is_empty(), "{}/{} available but no url", m.id, f.name);
468                assert_eq!(f.sha256.len(), 64, "{}/{} sha256 must be 64 hex chars", m.id, f.name);
469                assert!(
470                    f.sha256.bytes().all(|b| b.is_ascii_hexdigit()),
471                    "{}/{} sha256 not hex",
472                    m.id,
473                    f.name
474                );
475            }
476            assert!(m.is_available());
477        }
478    }
479}
qsm_core/models/registry.rs

qsm_core/models/
registry.rs
