FrancescoLR/OMNI-MS-model

OMNI-MS: contrast-agnostic multiple sclerosis lesion segmentation

OMNI-MS segments white-matter lesions (WML) in brain MRI. It couples a 3DINO ViT-L backbone (pretrained on UK Biobank brain MRI) with a ViT-Adapter + UNETR decoder, and is trained on real images mixed with synthetic images generated on the fly from label maps (random tissue contrast, lesion contrast, bias field, blur/resolution and noise), following the SynthSeg idea, so that one model can be applied across MRI contrasts.

Try it in the demo Space: https://huggingface.co/spaces/FrancescoLR/OMNI-MS

Files

File	Content
OMNI-MS.pth	Complete network state_dict (backbone + adapter + decoder, 350 M parameters, fp32, 1.4 GB). No other weights are needed.

Model details

• Architecture: 3DINO ViT-L (patch 16) + ViT-Adapter + UNETR decoder; 1 input channel, 2 output classes (background / WML).
• Training data: 931 subjects (837 train / 94 validation) mixing FLAIR (668), 3T T1-weighted (64), 7T T1-weighted (48) and MP2RAGE UNIT1 (60 + 91; the 91 are 7T) scans with lesion masks, plus synthetic images drawn from the associated label maps. The proportion of synthetic images starts at 80 % and decays linearly to 0 % over training (curriculum).
• Training: 200 epochs x 250 iterations, 112^3 voxel patches, DiceFocal loss, last 4 ViT blocks unfrozen.
• Inputs were skull-stripped during training. Whole-head images degrade performance markedly.

Inference

Use the code in the Space repository (omni_ms.py and the dinov2/ package):

import os
os.environ["XFORMERS_DISABLED"] = "1"
from huggingface_hub import hf_hub_download
import omni_ms

weights = hf_hub_download("FrancescoLR/OMNI-MS-model", "OMNI-MS.pth")
model = omni_ms.build_model(weights, "cuda")
omni_ms.segment(model, "flair_skullstripped.nii.gz", "lesion_mask.nii.gz", "cuda")

Settings used for all reported results: RAS reorientation; axes coarser than 1 mm resampled to 1 mm (finer axes kept native); foreground crop; z-score over non-zero voxels; 112^3 sliding window, 75 % overlap; argmax; prediction resampled back to the input grid (nearest neighbour). Requires a CUDA GPU (bf16 autocast); peak GPU memory was about 11 GB in tests, and inference took 8-31 s per image on an L40S.

Results (Dice, held-out data not used for training)

Test set	Contrast	n	Dice
Validation split of the training pool (own domain)	mixed	94	0.722
MSSEG-2, research	FLAIR	14	0.712
MSSEG-2, clinical (thick slices)	FLAIR	10	0.668
MSLesSeg	FLAIR	93	0.695
Sinai 7T longitudinal cohort	MP2RAGE UNIT1 (7T)	14	0.561
Sinai IRSWIET	T2*-weighted (IRSWIET)	14	0.490

Dice is voxel-wise on the foreground label, after skull-stripping where noted below. The Sinai 7T cohort overlaps by subject ID with 4 subjects of the 7T training data (same-ID scans may be the same participants); excluding them gives 0.574 on the remaining 10. The Sinai IRSWIET set was skull-stripped with SynthStrip (--no-csf); on whole-head images Dice drops to 0.392. Performance on contrasts other than FLAIR is clearly lower and varies between subjects; two of the 14 IRSWIET subjects fail almost completely.

Intended use and limitations

Research use only; not a medical device and not intended for clinical decision making. Lesion-level errors are common (see lesion-wise precision and recall in the paper once available), and small or confluent lesions are the hardest cases. Please visually check every output.

License

The model builds on 3DINO (CC BY-NC-ND 4.0, Sunnybrook Research Institute). OMNI-MS is therefore released for non-commercial use under the same terms; see the 3DINO repository for the full conditions.

Citation

OMNI-MS: reference to be added. Please also cite 3DINO:

@article{xu3dino2025,
  title={A generalizable 3D framework and model for self-supervised learning in medical imaging},
  author={Xu, Tony and Hosseini, Sepehr and Anderson, Chris and Rinaldi, Anthony and Krishnan, Rahul G. and Martel, Anne L. and Goubran, Maged},
  journal={npj Digital Medicine},
  year={2025},
  doi={10.1038/s41746-025-02035-w},
}