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`, or [`super::provide_weights`]).
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, NetworkMemory, WeightFile, WeightStatus};
16
17/// A [`NetworkMemory`], in the order its fields are listed: resident, load peak, inference fixed
18/// bytes and bytes per input voxel (native), WASM fixed bytes and bytes per input voxel. The
19/// figures come from `examples/inference_memory.rs` and the `browser_memory::net_*` cases in
20/// `tests/wasm_pipeline.rs` (see the [module docs](super#memory)).
21const fn mem(
22    resident_bytes: u64,
23    load_peak_bytes: u64,
24    inference_fixed_bytes: u64,
25    inference_bytes_per_voxel: u64,
26    wasm_fixed_bytes: u64,
27    wasm_bytes_per_voxel: u64,
28) -> NetworkMemory {
29    NetworkMemory {
30        resident_bytes,
31        load_peak_bytes,
32        inference_fixed_bytes,
33        inference_bytes_per_voxel,
34        wasm_fixed_bytes,
35        wasm_bytes_per_voxel,
36    }
37}
38
39/// All models known to QSM-Core, in a stable order.
40pub fn all_models() -> &'static [ModelSpec] {
41    MODELS
42}
43
44/// Look up a model by its [`ModelSpec::id`] (case-sensitive), e.g. `"qsmnet"`.
45pub fn find_model(id: &str) -> Option<&'static ModelSpec> {
46    MODELS.iter().find(|m| m.id == id)
47}
48
49// A single ONNX weight file whose location is not yet known (Pending model).
50const fn pending_onnx(name: &'static str) -> WeightFile {
51    WeightFile { name, url: "", sha256: "", bytes: 0 }
52}
53
54const MODELS: &[ModelSpec] = &[
55    // ---- Background field removal ------------------------------------------
56    ModelSpec {
57        id: "bfrnet",
58        name: "BFRnet",
59        stage: ModelStage::BackgroundRemoval,
60        status: WeightStatus::Available,
61        origin: Framework::Matlab,
62        description: "Dual-frequency octave-convolution U-Net for background field \
63                      removal (total field → local field). Fully convolutional.",
64        paper: "Zhu et al., Z Med Phys 2023; doi:10.1016/j.zemedi.2022.08.001",
65        source: "https://github.com/sunhongfu/BFRnet",
66        license: "",
67        // Mirrored on Hugging Face (qsmxt/qsm-onnx-weights). Verified: anonymous
68        // download + SHA-256 match.
69        files: &[WeightFile {
70            name: "bfrnet.onnx",
71            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/bfrnet.onnx",
72            sha256: "6f693f0a02c94550179c4b5188ce652fc4bd8198ff57aaddd102fba67fe873d7",
73            bytes: 79_600_612,
74        }],
75        inputs: &["field"],
76        outputs: &["local_field"],
77        size_divisor: 8,
78        // Measured in WASM: the heap's slope over 48³–80³ patches (`browser_memory::tile_cost_*`).
79        tile_bytes_per_voxel: Some(2_054),
80        memory: mem(82_100_000, 164_200_000, 128_000_000, 1738, 245_000_000, 2556),
81    },
82    // ---- Dipole inversion --------------------------------------------------
83    ModelSpec {
84        id: "xqsm",
85        name: "xQSM",
86        stage: ModelStage::DipoleInversion,
87        status: WeightStatus::Available,
88        origin: Framework::PyTorch,
89        description: "Octave-convolution U-Net with a learned residual for dipole \
90                      inversion (local field → susceptibility).",
91        paper: "Gao et al., NMR Biomed 2021; doi:10.1002/nbm.4461",
92        source: "https://github.com/sunhongfu/xQSM",
93        license: "",
94        // Exported from xQSM_invivo.pth (v1.0-demo) with scripts/onnx-export/export_xqsm.py;
95        // torch↔onnxruntime parity max|Δ| ≈ 6e-5. Mirrored on Hugging Face (qsmxt/qsm-onnx-weights).
96        files: &[WeightFile {
97            name: "xqsm.onnx",
98            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/xqsm.onnx",
99            sha256: "81854ec2ca85bba25c2f9aae05efdcac299d0efffb03fe6a0e6797019e46b487",
100            bytes: 20_901_826,
101        }],
102        inputs: &["field"],
103        outputs: &["chi"],
104        size_divisor: 8,
105        // Measured in WASM: the heap's slope over 48³–80³ patches (`browser_memory::tile_cost_*`).
106        tile_bytes_per_voxel: Some(1_146),
107        memory: mem(21_600_000, 43_800_000, 59_000_000, 913, 243_000_000, 980),
108    },
109    ModelSpec {
110        id: "qsmnet",
111        name: "QSMnet",
112        stage: ModelStage::DipoleInversion,
113        status: WeightStatus::Available,
114        origin: Framework::TensorFlow,
115        description: "3D U-Net dipole inversion trained on COSMOS. Expects 1 mm \
116                      isotropic input; z-scored with training mean/std.",
117        paper: "Yoon et al., NeuroImage 2018; doi:10.1016/j.neuroimage.2018.06.030",
118        source: "https://github.com/SNU-LIST/QSMnet",
119        license: "",
120        // Clean PyTorch re-export of the TF1.14 checkpoint (tract-friendly NCDHW), with the
121        // checkpoint's own ReLU activation (QSM.rs#330); see scripts/onnx-export/export_qsmnet.py.
122        // Mirrored on Hugging Face (qsmxt/qsm-onnx-weights).
123        files: &[WeightFile {
124            name: "qsmnet.onnx",
125            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/qsmnet.onnx",
126            sha256: "f1b625e3c04e679077cb8ac9b2617b14e4bbe452f9ae516a806a1acd36e2b159",
127            bytes: 397_801_168,
128        }],
129        inputs: &["field"],
130        outputs: &["chi"],
131        size_divisor: 16,
132        // Measured in WASM over 48³–80³ patches: QSMnet 1 680, QSMnet+ 1 694; the one tiled wrapper
133        // serves both, so both carry the larger.
134        tile_bytes_per_voxel: Some(1_694),
135        memory: mem(398_000_000, 796_200_000, 443_000_000, 842, 1_008_000_000, 1068),
136    },
137    ModelSpec {
138        id: "qsmnet-plus",
139        name: "QSMnet+",
140        stage: ModelStage::DipoleInversion,
141        status: WeightStatus::Available,
142        origin: Framework::TensorFlow,
143        description: "QSMnet retrained with susceptibility-scaling augmentation for \
144                      a wider, more linear χ range.",
145        paper: "Jung et al., NeuroImage 2020; doi:10.1016/j.neuroimage.2020.116579",
146        source: "https://github.com/SNU-LIST/QSMnet",
147        license: "",
148        // Clean PyTorch re-export of the TF1.14 QSMnet+_64 checkpoint (same U-Net
149        // as QSMnet, different weights + norm). Mirrored on Hugging Face (qsmxt/qsm-onnx-weights).
150        files: &[WeightFile {
151            name: "qsmnet-plus.onnx",
152            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/qsmnet-plus.onnx",
153            sha256: "0ed31f9a1b66f75fee4a96022b6bc3838b91bc5b819714dade5ca87cd331683d",
154            bytes: 397_801_854,
155        }],
156        inputs: &["field"],
157        outputs: &["chi"],
158        size_divisor: 16,
159        // Measured in WASM over 48³–80³ patches: QSMnet 1 680, QSMnet+ 1 694; the one tiled wrapper
160        // serves both, so both carry the larger.
161        tile_bytes_per_voxel: Some(1_694),
162        memory: mem(398_000_000, 796_200_000, 443_000_000, 842, 1_008_000_000, 1068),
163    },
164    ModelSpec {
165        id: "qsmnetme",
166        name: "QSMnetME",
167        stage: ModelStage::DipoleInversion,
168        status: WeightStatus::Available,
169        origin: Framework::Onnx,
170        description: "The SNU-LIST χ-separation toolbox's QSM network (already ONNX): a \
171                      QSMnet-style 3D U-Net trained on multi-echo data, run by the toolbox's \
172                      route (field rotated to B0 along z and resampled to 1 mm, z-scored with \
173                      the χ-sepnet field statistics, whole volume, mapped back). Spatial axes \
174                      re-declared dynamic (the toolbox ships a fixed 192×192×128 input). \
175                      Normalization constants are baked into the Rust glue (QsmnetmeNorm).",
176        paper: "Kim et al., Hum Brain Mapp 2025 (doi:10.1002/hbm.70136)",
177        source: "https://github.com/SNU-LIST/chi-separation",
178        license: "",
179        // χ-separation toolbox v1.1.3 `models/240904_QSMnet.onnx` (SHA-256 de2f2c23…) with its
180        // spatial axes re-declared dynamic by scripts/onnx-export/redeclare_dynamic_axes.py:
181        // identical output at the declared 192×192×128. Mirrored on Hugging Face
182        // (qsmxt/qsm-onnx-weights, commit e7116c8); anonymous download + SHA-256 checked.
183        files: &[WeightFile {
184            name: "qsmnetme.onnx",
185            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/qsmnetme.onnx",
186            sha256: "c2f030b625b2fc83f3ae45b38c996e006a069444d2da32d2b6022045c92c01e5",
187            bytes: 90_307_128,
188        }],
189        inputs: &["field"],
190        outputs: &["chi"],
191        size_divisor: 16,
192        tile_bytes_per_voxel: None,
193        // The same graph as R2PRIMEnet's (same layers, shapes and byte count; other weights), so
194        // its measured memory figures.
195        memory: mem(90_500_000, 181_200_000, 114_000_000, 825, 188_000_000, 1155),
196    },
197    ModelSpec {
198        id: "qsmgan",
199        name: "QSMGAN",
200        stage: ModelStage::DipoleInversion,
201        status: WeightStatus::Available,
202        origin: Framework::PyTorch,
203        description: "3D U-Net generator (WGAN-GP refined) for dipole inversion, run \
204                      patch-wise (64³ input → 48³ output, i64o48). Only the generator is \
205                      used at inference; the sign flip, input_scale/tanh (χ=atanh/10) and \
206                      patch tiling live in the Rust glue.",
207        paper: "Chen et al., NeuroImage 2020; doi:10.1016/j.neuroimage.2019.116389",
208        source: "https://github.com/mmorri10/QSMGAN-LupoLab",
209        license: "MIT (fork)",
210        files: &[WeightFile {
211            name: "qsmgan.onnx",
212            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/qsmgan.onnx",
213            sha256: "238573a6fa8a563004fc28e63a30694289d1a1a204f3c406daba486b1122d604",
214            bytes: 10_753_955,
215        }],
216        inputs: &["field"],
217        outputs: &["chi"],
218        size_divisor: 8,
219        tile_bytes_per_voxel: None,
220        memory: mem(11_000_000, 22_200_000, 0, 518, 0, 699),
221    },
222    ModelSpec {
223        id: "lpcnn",
224        name: "LPCNN",
225        stage: ModelStage::DipoleInversion,
226        status: WeightStatus::Available,
227        origin: Framework::PyTorch,
228        description: "Learned proximal CNN, 3 unrolled iterations of proximal \
229                      gradient descent. The k-space dipole data-consistency step, the \
230                      unroll, the learned step size and the mean/std normalization run \
231                      in Rust (rustfft); only the learned proximal CNN is ONNX.",
232        paper: "Lai et al., MICCAI 2020; doi:10.1007/978-3-030-59713-9_13",
233        source: "https://github.com/Sulam-Group/LPCNN",
234        license: "",
235        files: &[WeightFile {
236            name: "lpcnn.onnx",
237            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/lpcnn.onnx",
238            sha256: "831f2a161f6ab8ee7bf64cc2dd8bca532fc971d0083cf5729d6302d3ed31625d",
239            bytes: 1_790_869,
240        }],
241        inputs: &["field", "mask", "b_vec"],
242        outputs: &["chi"],
243        size_divisor: 8,
244        // Measured in WASM: a whole 64³ tile per patch voxel, since this wrapper loads its network
245        // inside every patch, so each tile in flight also holds a parsed model.
246        tile_bytes_per_voxel: Some(964),
247        memory: mem(2_000_000, 4_100_000, 11_000_000, 661, 0, 766),
248    },
249    ModelSpec {
250        id: "ir2qsm",
251        name: "IR2QSM",
252        stage: ModelStage::DipoleInversion,
253        status: WeightStatus::Available,
254        origin: Framework::PyTorch,
255        description: "IR2U-net dipole inversion: a 3D U-net (depth 4) run for 4 \
256                      unrolled iterations with reverse concatenations and a recurrent \
257                      middle module. The whole net is ONNX; the Rust glue only does \
258                      /8 zero-pad, crop and mask (no normalization, ppm in/out). The \
259                      ungated inference-time AddNoise is pinned to its noise-free \
260                      branch for deterministic export.",
261        paper: "Li et al., Med. Phys. 2025; doi:10.1002/mp.17747; arXiv:2406.12300",
262        source: "https://github.com/YangGaoUQ/IR2QSM",
263        license: "",
264        files: &[WeightFile {
265            name: "ir2qsm.onnx",
266            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/ir2qsm.onnx",
267            sha256: "dcaf7a26d6633900f339f94be16d580ccea6147bbf6da26eacb9a5e1acdc0c53",
268            bytes: 40_259_354,
269        }],
270        inputs: &["field", "mask"],
271        outputs: &["chi"],
272        size_divisor: 8,
273        // Measured in WASM: the heap's slope over 48³–80³ patches (`browser_memory::tile_cost_*`).
274        tile_bytes_per_voxel: Some(904),
275        memory: mem(41_000_000, 82_900_000, 74_000_000, 704, 147_000_000, 769),
276    },
277    // ---- Single-step (phase/total field → χ) -------------------------------
278    ModelSpec {
279        id: "autoqsm",
280        name: "AutoQSM",
281        stage: ModelStage::SingleStep,
282        status: WeightStatus::Available,
283        origin: Framework::TensorFlow,
284        description: "V-Net single-step reconstruction (total field → susceptibility) \
285                      with no separate brain extraction. Fixed 64³→32³ patch net; \
286                      the Rust glue does the sliding-window tiling + blend.",
287        paper: "Wei et al., NeuroImage 2019; doi:10.1016/j.neuroimage.2019.116064",
288        source: "https://github.com/AMRI-Lab/AutoQSM",
289        license: "",
290        // Clean PyTorch re-export of the Keras V-Net (tract-friendly NCDHW);
291        // see scripts/onnx-export/export_autoqsm.py. Mirrored on Hugging Face (qsmxt/qsm-onnx-weights).
292        files: &[WeightFile {
293            name: "autoqsm.onnx",
294            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/autoqsm.onnx",
295            sha256: "c2f23c6ae735fd1677e8623da22cb6f25b4a932ae50ba1ef08a75439e240977e",
296            bytes: 5_609_662,
297        }],
298        // Fixed-size patch net; `size_divisor` is not a whole-volume constraint here
299        // (tiling handles padding), left at 1.
300        inputs: &["field"],
301        outputs: &["chi"],
302        size_divisor: 1,
303        tile_bytes_per_voxel: None,
304        memory: mem(5_800_000, 11_800_000, 0, 495, 0, 638),
305    },
306    ModelSpec {
307        id: "iqsm",
308        name: "iQSM",
309        stage: ModelStage::SingleStep,
310        status: WeightStatus::Available,
311        origin: Framework::PyTorch,
312        description: "LoT-Unet single-step reconstruction from wrapped phase to \
313                      susceptibility. The learnable-Laplacian front-end is fused into \
314                      the exported graph; inputs are phase, mask, TE (s), B0 (T).",
315        paper: "Gao et al., NeuroImage 2022; doi:10.1016/j.neuroimage.2022.119410",
316        source: "https://github.com/sunhongfu/iQSM",
317        license: "",
318        files: &[WeightFile {
319            name: "iqsm.onnx",
320            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/iqsm.onnx",
321            sha256: "8538a33892a877812e6cd0e22927a5040827c126e9347073d3dd0934bf226b09",
322            bytes: 17_233_763,
323        }],
324        inputs: &["phase", "mask", "te", "b0", "border"],
325        outputs: &["chi"],
326        size_divisor: 16,
327        // Measured in WASM: the heap's slope over 48³–80³ patches (`browser_memory::tile_cost_*`).
328        tile_bytes_per_voxel: Some(600),
329        memory: mem(17_600_000, 35_300_000, 34_000_000, 476, 55_000_000, 571),
330    },
331    ModelSpec {
332        id: "iqsm-plus",
333        name: "iQSM+",
334        stage: ModelStage::SingleStep,
335        status: WeightStatus::Available,
336        origin: Framework::PyTorch,
337        description: "iQSM with orientation-adaptive latent feature editing (OA-LFE); \
338                      the B0 direction is a genuine network input. Inputs: phase, \
339                      mask, TE, B0, z_prjs (B0 dir), border.",
340        paper: "Gao et al., Med Image Anal 2024; doi:10.1016/j.media.2024.103160",
341        source: "https://github.com/sunhongfu/iQSM_Plus",
342        license: "",
343        files: &[WeightFile {
344            name: "iqsm-plus.onnx",
345            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/iqsm-plus.onnx",
346            sha256: "f93a4c83804759a4a0d24603f7754efb997ddd2ba5e3a0346ed6493725cafb1d",
347            bytes: 17_740_526,
348        }],
349        inputs: &["phase", "mask", "te", "b0", "z_prjs", "border"],
350        outputs: &["chi"],
351        size_divisor: 16,
352        // Measured in WASM: the heap's slope over 48³–80³ patches (`browser_memory::tile_cost_*`).
353        tile_bytes_per_voxel: Some(1_149),
354        memory: mem(20_300_000, 41_000_000, 31_000_000, 1055, 65_000_000, 1037),
355    },
356    ModelSpec {
357        id: "iqfm",
358        name: "iQFM",
359        stage: ModelStage::PhaseToField,
360        status: WeightStatus::Available,
361        origin: Framework::PyTorch,
362        description: "The tissue-field head of the iQSM LoT-Unet: wrapped phase → \
363                      local (background-removed) field in one network (joint unwrap + \
364                      BFR). Same architecture/inputs as iQSM, `lfs` weights; output is \
365                      the local field (ppm), not susceptibility.",
366        paper: "Gao et al., NeuroImage 2022; doi:10.1016/j.neuroimage.2022.119410",
367        source: "https://github.com/sunhongfu/iQSM",
368        license: "",
369        files: &[WeightFile {
370            name: "iqfm.onnx",
371            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/iqfm.onnx",
372            sha256: "81965f4c0981612d979977c9618847e4cb1d340efa287b8dde894bdea5bf6ac9",
373            bytes: 17_233_763,
374        }],
375        inputs: &["phase", "mask", "te", "b0", "border"],
376        outputs: &["localfield"],
377        size_divisor: 16,
378        // Measured in WASM: the heap's slope over 48³–80³ patches (`browser_memory::tile_cost_*`).
379        tile_bytes_per_voxel: Some(600),
380        memory: mem(17_600_000, 35_300_000, 34_000_000, 476, 55_000_000, 571),
381    },
382    ModelSpec {
383        id: "nextqsm",
384        name: "NeXtQSM",
385        stage: ModelStage::SingleStep,
386        status: WeightStatus::Available,
387        origin: Framework::TensorFlow,
388        description: "U-Net background removal followed by a 6-step variational-network \
389                      dipole inversion (total field → susceptibility). Reimplemented as \
390                      a Rust hybrid: two exported U-Nets (`nextqsm-bf` = BFR forward, \
391                      `nextqsm-vjp` = the regularizer gradient ∇ₓ mean|VarNet(x)| written \
392                      as a forward graph) plus the FFT data-consistency gradient and unroll \
393                      in Rust. Order the files BFR-first, VJP-second.",
394        paper: "Cognolato et al., Medical Image Analysis 84:102700 (2023); doi:10.1016/j.media.2022.102700",
395        source: "https://github.com/QSMxT/NeXtQSM",
396        license: "MIT",
397        files: &[
398            WeightFile {
399                name: "nextqsm-bf.onnx",
400            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/nextqsm-bf.onnx",
401                sha256: "d0f0b0391153f8e5e07ef85d46fbe492afc8394fab008350540a6868d0b4a3ba",
402                bytes: 113_211_474,
403            },
404            WeightFile {
405                name: "nextqsm-vjp.onnx",
406            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/nextqsm-vjp.onnx",
407                sha256: "50fe388663c05b1d3996139975270e500e13feb025f1b792576dbc0da814a4bc",
408                bytes: 42_437_407,
409            },
410        ],
411        inputs: &["field", "mask", "b_vec"],
412        outputs: &["chi"],
413        size_divisor: 64,
414        // Measured in WASM: a whole 64³ tile per patch voxel, since this wrapper loads its network
415        // inside every patch, so each tile in flight also holds a parsed model.
416        tile_bytes_per_voxel: Some(1_952),
417        memory: mem(156_200_000, 226_900_000, 110_000_000, 600, 316_000_000, 753),
418    },
419    ModelSpec {
420        id: "modl-qsm",
421        name: "MoDL-QSM",
422        stage: ModelStage::DipoleInversion,
423        status: WeightStatus::Available,
424        origin: Framework::TensorFlow,
425        description: "Model-based deep learning: 3-iteration unroll of a QSM \
426                      forward-model gradient descent with a learned 2-channel CNN \
427                      prior. The k-space dipole data-consistency (A/Aᴴ), the unroll, \
428                      the learned step size and per-channel mean/std normalization run \
429                      in Rust (rustfft); only the CNN prior is ONNX. Output is the STI \
430                      χ33 component.",
431        paper: "Feng et al., NeuroImage 2021; doi:10.1016/j.neuroimage.2021.118376",
432        source: "https://github.com/Ruimin-Feng/MoDL-QSM",
433        license: "",
434        files: &[WeightFile {
435            name: "modl-qsm.onnx",
436            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/modl-qsm.onnx",
437            sha256: "1f8da7d39eaaa4c3f0f9e67298fba5f14230e769cfd6349c81d0b74a6f800c57",
438            bytes: 1_794_143,
439        }],
440        inputs: &["field", "mask", "b_vec"],
441        outputs: &["chi"],
442        size_divisor: 1,
443        // Measured in WASM: a whole 64³ tile per patch voxel, since this wrapper loads its network
444        // inside every patch, so each tile in flight also holds a parsed model.
445        tile_bytes_per_voxel: Some(909),
446        memory: mem(2_000_000, 4_100_000, 11_000_000, 665, 12_000_000, 629),
447    },
448    // ---- χ-separation (χ+ / χ−) --------------------------------------------
449    ModelSpec {
450        id: "susep-net",
451        name: "SUSEP-Net",
452        stage: ModelStage::ChiSeparation,
453        status: WeightStatus::Available,
454        origin: Framework::PyTorch,
455        description: "Dual-branch 3D U-Net source separation from z-scored \
456                      [QSM, R2', local field] → [χ+, χ−]. Normalization constants \
457                      are baked into the Rust glue (SusepNetNorm).",
458        paper: "Li, Gao, Sun et al., arXiv:2506.13293 (2025)",
459        source: "https://github.com/YangGaoUQ/SUSEP-Net",
460        license: "",
461        files: &[WeightFile {
462            name: "susep-net.onnx",
463            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/susep-net.onnx",
464            sha256: "9bc85aa28451fda2c661de4b8fae8604ced7a54f028f94ed33629f6111a48b03",
465            bytes: 205_743_584,
466        }],
467        inputs: &["qsm", "r2prime", "lfs"],
468        outputs: &["chi_pos", "chi_neg"],
469        size_divisor: 8,
470        tile_bytes_per_voxel: None,
471        memory: mem(206_100_000, 412_700_000, 249_000_000, 1960, 553_000_000, 2164),
472    },
473    ModelSpec {
474        id: "r2primenet",
475        name: "R2PRIMEnet",
476        stage: ModelStage::R2PrimeGeneration,
477        status: WeightStatus::Available,
478        origin: Framework::Onnx,
479        description: "SNU-LIST R2*→R2′ conversion network from the χ-sepnet pipeline \
480                      (already ONNX): a fully-convolutional 3D U-Net mapping the Dr-scaled, \
481                      z-scored R2* map to R2′, run as an overlapping sliding window. Supplies \
482                      R2′ for the GRE-only condition, where no spin-echo R2 is measured. \
483                      Spatial axes re-declared dynamic (the authors ship a fixed 192×192×128 \
484                      input) so 32-bit hosts can use a smaller patch. Normalization constants \
485                      (Dr=114) are baked into the Rust glue (R2PrimeNetNorm).",
486        paper: "Kim et al., Hum Brain Mapp 2025 (doi:10.1002/hbm.70136)",
487        source: "https://github.com/SNU-LIST/chi_sepnet",
488        license: "",
489        files: &[WeightFile {
490            name: "r2primenet.onnx",
491            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/r2primenet.onnx",
492            sha256: "44cb5e67d1c68dae87a5f532e501dc0d4bf627d91f9c8ba35a56df5ac7ec3cd0",
493            bytes: 90_307_128,
494        }],
495        inputs: &["r2star"],
496        outputs: &["r2prime"],
497        size_divisor: 16,
498        tile_bytes_per_voxel: None,
499        memory: mem(90_500_000, 181_200_000, 114_000_000, 825, 188_000_000, 1155),
500    },
501    ModelSpec {
502        id: "chi-sepnet",
503        name: "χ-sepnet",
504        stage: ModelStage::ChiSeparation,
505        status: WeightStatus::Available,
506        origin: Framework::Onnx,
507        description: "SNU-LIST χ-separation network (already ONNX): a fully-convolutional \
508                      3D U-Net mapping [QSM, local field, R2′/Dr] (z-scored) → [χ+, χ−], \
509                      run by the toolbox's route (training frame, QSMnetME QSM channel, whole \
510                      volume) or alone as an overlapping sliding window. Spatial axes re-declared dynamic \
511                      (the authors ship a fixed 192×192×128 input) so 32-bit hosts can use a \
512                      smaller patch. Normalization constants (Dr=114) are baked into the Rust \
513                      glue (ChiSepNetNorm).",
514        paper: "Kim et al. / SNU-LIST chi-separation toolbox",
515        source: "https://github.com/SNU-LIST/chi-separation",
516        license: "",
517        files: &[WeightFile {
518            name: "chi-sepnet.onnx",
519            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/chi-sepnet.onnx",
520            sha256: "5b442fdfdb88f50ec9149b384dabd0d2d9adb6983947d9cfa58382b78676bed9",
521            bytes: 90_314_172,
522        }],
523        inputs: &["local_field", "qsm", "r2prime"],
524        outputs: &["chi_pos", "chi_neg"],
525        size_divisor: 8,
526        tile_bytes_per_voxel: None,
527        memory: mem(90_500_000, 181_200_000, 114_000_000, 833, 188_000_000, 1155),
528    },
529    ModelSpec {
530        id: "hd-bet",
531        name: "HD-BET",
532        stage: ModelStage::BrainExtraction,
533        status: WeightStatus::Available,
534        origin: Framework::PyTorch,
535        description: "nnU-Net v2 3D U-Net brain extraction (HD-BET v2): magnitude → brain \
536                      mask, trained on 11,751 multi-sequence clinical MRIs. Runs at 1 mm \
537                      on the whole volume (or Gaussian-weighted sliding-window patches, \
538                      96×192×192 as trained); nnU-Net's crop/normalise/resample pipeline \
539                      is in the Rust glue (bet::hd_bet). Spatial axes are dynamic \
540                      (multiples of 16×32×32).",
541        paper: "Isensee et al., Hum Brain Mapp 40(17):4952-4964 (2019); https://doi.org/10.1002/hbm.24750",
542        source: "https://github.com/MIC-DKFZ/HD-BET",
543        license: "CC-BY-NC-4.0",
544        files: &[WeightFile {
545            name: "hd-bet.onnx",
546            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/hd-bet.onnx",
547            sha256: "f15200f9a0697cf53b151e316a0c4fe2cacfdbe2220deb36cb56184be4a227f9",
548            bytes: 123_168_028,
549        }],
550        inputs: &["image"],
551        outputs: &["logits"],
552        size_divisor: 32,
553        tile_bytes_per_voxel: None,
554        memory: mem(123_500_000, 247_200_000, 162_000_000, 821, 307_000_000, 1101),
555    },
556    ModelSpec {
557        id: "rs2-net",
558        name: "RS2-Net",
559        stage: ModelStage::BrainExtraction,
560        status: WeightStatus::Available,
561        origin: Framework::PyTorch,
562        description: "Swin-UNETR rodent (mouse, rat) brain extraction, trained within nnU-Net \
563                      v2 on 1,142 MRIs from 89 centres: magnitude → brain mask. Runs at \
564                      0.25×0.2×0.16 mm with Gaussian-weighted sliding-window patches; nnU-Net's \
565                      transpose/crop/normalise/resample pipeline is in the Rust glue \
566                      (bet::rs2_net). The graph is traced at a fixed 128×96×128 patch.",
567        paper: "Lin et al., NeuroImage 298:120769 (2024); https://doi.org/10.1016/j.neuroimage.2024.120769",
568        source: "https://github.com/VitoLin21/Rodent-Skull-Stripping",
569        license: "GPL-3.0",
570        // Exported from the released RS2_pretrained_model.pt with
571        // scripts/onnx-export/export_rs2net.py (reproducible: identical bytes on re-export);
572        // torch↔onnxruntime sign agreement 99.9998 %.
573        files: &[WeightFile {
574            name: "rs2-net.onnx",
575            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/rs2-net.onnx",
576            sha256: "a120ce43b06a9f3ddecf85bf18f281cbb54c354ef15b34a6bcaaffa4e3c4a570",
577            bytes: 63_456_162,
578        }],
579        inputs: &["input"],
580        outputs: &["logits"],
581        size_divisor: 32,
582        tile_bytes_per_voxel: None,
583        memory: mem(83_100_000, 159_600_000, 0, 1692, 0, 1697),
584    },
585    ModelSpec {
586        id: "mindgrab",
587        name: "MindGrab",
588        stage: ModelStage::BrainExtraction,
589        status: WeightStatus::Available,
590        origin: Framework::PyTorch,
591        description: "MeshNet brain extraction for any contrast (25 dilated 3×3×3 convolutions, \
592                      15 channels, 146 k parameters) trained only on synthetic data: image → \
593                      brain mask. Runs on the whole 256³ 1 mm LIA conformed volume; the \
594                      conform/normalise/largest-component/reslice pipeline is in the Rust glue \
595                      (bet::mindgrab). Spatial axes are dynamic.",
596        paper: "Fani et al., arXiv:2506.11860 (2025); https://doi.org/10.48550/arXiv.2506.11860",
597        source: "https://github.com/neuroneural/brainchop-cli",
598        license: "MIT",
599        // Exported from brainchop-models meshnet/mindgrab (model.json + model.pth) with
600        // scripts/onnx-export/export_mindgrab.py. Mirrored on Hugging Face
601        // (qsmxt/qsm-onnx-weights). Verified: anonymous download + SHA-256 match.
602        files: &[WeightFile {
603            name: "mindgrab.onnx",
604            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/mindgrab.onnx",
605            sha256: "2e20e5e03211976d4c30b4e704c57f7f18cb1b78e840105acf12765c6e3c342b",
606            bytes: 672_150,
607        }],
608        inputs: &["image"],
609        outputs: &["logits"],
610        size_divisor: 1,
611        tile_bytes_per_voxel: None,
612        memory: mem(2_000_000, 3_600_000, 206_000_000, 426, 363_000_000, 311),
613    },
614    // ---- Anatomical segmentation -------------------------------------------
615    ModelSpec {
616        id: "synthseg",
617        name: "SynthSeg 1.0",
618        stage: ModelStage::Segmentation,
619        status: WeightStatus::Available,
620        origin: Framework::TensorFlow,
621        description: "Contrast-agnostic whole-brain segmentation: magnitude → 32 FreeSurfer \
622                      labels. A 5-level 3D U-Net (24 features, ELU, 13.2 M parameters) trained \
623                      only on synthetic images with randomised contrast, so it runs directly on \
624                      GRE magnitude without a T1w scan. Resample/orient/normalise, flip \
625                      averaging and the topological cleanup are in the Rust glue \
626                      (segment::synthseg). Spatial axes are dynamic (multiples of 32).",
627        paper: "Billot et al., Med Image Anal 86:102789 (2023); https://doi.org/10.1016/j.media.2023.102789",
628        source: "https://github.com/BBillot/SynthSeg",
629        license: "Apache-2.0",
630        // Mirrored on Hugging Face (qsmxt/qsm-onnx-weights). Verified: anonymous
631        // download + SHA-256 match.
632        files: &[WeightFile {
633            name: "synthseg.onnx",
634            url: "https://huggingface.co/qsmxt/qsm-onnx-weights/resolve/main/synthseg.onnx",
635            sha256: "c2821a74e8a03d4073896b5c2e359b3ef86f9776bacd78e00da1c757aa97bbcc",
636            bytes: 52_998_326,
637        }],
638        inputs: &["image"],
639        outputs: &["unet_prediction"],
640        size_divisor: 32,
641        tile_bytes_per_voxel: None,
642        memory: mem(53_400_000, 106_800_000, 87_000_000, 919, 132_000_000, 1187),
643    },
644    ModelSpec {
645        id: "synthseg-2.0",
646        name: "SynthSeg 2.0",
647        stage: ModelStage::Segmentation,
648        status: WeightStatus::Pending,
649        origin: Framework::TensorFlow,
650        description: "SynthSeg 2.0: same architecture as `synthseg` with a 33-label set (adds \
651                      a general CSF class). The upstream weights are not in the SynthSeg \
652                      repository — they ship with FreeSurfer, or come from the UCL download \
653                      linked in the SynthSeg README — so they have to be fetched before \
654                      exporting. Run it with SynthSegVersion::V2.",
655        paper: "Billot et al., PNAS 120(9):e2216399120 (2023); https://doi.org/10.1073/pnas.2216399120",
656        source: "https://github.com/BBillot/SynthSeg",
657        license: "Apache-2.0",
658        files: &[pending_onnx("synthseg-2.0.onnx")],
659        inputs: &["image"],
660        outputs: &["unet_prediction"],
661        size_divisor: 32,
662        tile_bytes_per_voxel: None,
663        // Not measured (the weights are not hosted): SynthSeg 1.0's figures, plus its 33rd output
664        // channel (the output tensor and its copies, about 12 B per voxel natively, 16 in WASM).
665        memory: mem(53_400_000, 106_800_000, 87_000_000, 931, 132_000_000, 1203),
666    },
667];
668
669#[cfg(test)]
670mod tests {
671    use super::*;
672
673    #[test]
674    fn ids_are_unique_and_findable() {
675        let mut seen = std::collections::HashSet::new();
676        for m in MODELS {
677            assert!(seen.insert(m.id), "duplicate model id: {}", m.id);
678            assert!(find_model(m.id).is_some(), "find_model failed for {}", m.id);
679            assert!(!m.name.is_empty(), "{} has empty name", m.id);
680            assert!(!m.files.is_empty(), "{} has no weight files", m.id);
681        }
682        assert!(find_model("does-not-exist").is_none());
683    }
684
685    #[test]
686    fn pending_models_are_not_available() {
687        // A Pending model is not runnable via download: it must lack a URL on at
688        // least one file (a known SHA-256 ahead of hosting is fine, e.g. BFRnet).
689        for m in MODELS.iter().filter(|m| m.status == WeightStatus::Pending) {
690            assert!(!m.is_available(), "{} is Pending but reports available", m.id);
691            assert!(
692                m.files.iter().any(|f| f.url.is_empty()),
693                "{} is Pending but every file has a url — flip status to Available",
694                m.id
695            );
696            // Any precomputed hash must still be a valid 64-char hex digest.
697            for f in m.files.iter().filter(|f| !f.sha256.is_empty()) {
698                assert_eq!(f.sha256.len(), 64, "{}/{} sha256 must be 64 hex chars", m.id, f.name);
699                assert!(
700                    f.sha256.bytes().all(|b| b.is_ascii_hexdigit()),
701                    "{}/{} sha256 not hex",
702                    m.id,
703                    f.name
704                );
705            }
706        }
707    }
708
709    #[test]
710    fn available_models_are_fully_specified() {
711        // Invariant: anything marked Available must have url + hex sha256 per file.
712        for m in MODELS.iter().filter(|m| m.status == WeightStatus::Available) {
713            for f in m.files {
714                assert!(!f.url.is_empty(), "{}/{} available but no url", m.id, f.name);
715                assert_eq!(f.sha256.len(), 64, "{}/{} sha256 must be 64 hex chars", m.id, f.name);
716                assert!(
717                    f.sha256.bytes().all(|b| b.is_ascii_hexdigit()),
718                    "{}/{} sha256 not hex",
719                    m.id,
720                    f.name
721                );
722            }
723            assert!(m.is_available());
724        }
725    }
726}
qsm_core/models/registry.rs

qsm_core/models/
registry.rs
