| # MS3SEG: Pre-trained Models for MS Lesion Segmentation |
|
|
| [](https://doi.org/10.6084/m9.figshare.30393475) |
| [](https://doi.org/10.6084/m9.figshare.30393475) |
| [](https://github.com/Mahdi-Bashiri/MS3SEG) |
| [](https://creativecommons.org/licenses/by/4.0/) |
|
|
| Pre-trained deep learning models for Multiple Sclerosis lesion segmentation from the **MS3SEG dataset**. |
|
|
| > **Note:** These are representative models from Fold 4 of our 5-fold cross-validation. Complete training code and all fold results are available in our [GitHub repository](https://github.com/Mahdi-Bashiri/MS3SEG). |
|
|
| --- |
|
|
| ## 📋 Repository Contents |
|
|
| ``` |
| MS3SEG/ |
| ├── kfold_brain_segmentation_20250924_232752_unified_focal_loss/models/ |
| │ ├── binary_abnormal_wmh/ # Binary MS lesion segmentation |
| │ │ ├── u-net_fold_4_best.h5 |
| │ │ ├── unet++_fold_4_best.h5 |
| │ │ ├── unetr_fold_4_best.h5 |
| │ │ └── swinunetr_fold_4_best.h5 |
| │ │ |
| │ ├── binary_ventricles/ # Binary ventricle segmentation |
| │ │ ├── u-net_fold_4_best.h5 |
| │ │ ├── unet++_fold_4_best.h5 |
| │ │ ├── unetr_fold_4_best.h5 |
| │ │ └── swinunetr_fold_4_best.h5 |
| │ │ |
| │ └── multi_class/ # 4-class tri-mask segmentation |
| │ │ ├── u-net_fold_4_best.h5 |
| │ │ ├── unet++_fold_4_best.h5 |
| │ │ ├── unetr_fold_4_best.h5 |
| │ │ └── swinunetr_fold_4_best.h5 |
| │ |
| ├── figures/ |
| │ ├── training_curves/ # Loss and metrics across epochs |
| │ └── sample_predictions/ # Visual results from paper |
| │ |
| ├── config/ |
| │ └── experiment_config.json # Model training configuration |
| └── README.md # This file |
| ``` |
|
|
| **Total Size:** ~1.2 GB (12 model files) |
|
|
| --- |
|
|
| ## 🎯 Model Overview |
|
|
| ### Segmentation Scenarios |
|
|
| | Scenario | Classes | Description | |
| |----------|---------|-------------| |
| | **Multi-class** | 4 | Background, Ventricles, Normal WMH, Abnormal WMH (MS lesions) | |
| | **Binary Lesion** | 2 | MS lesions vs. everything else | |
| | **Binary Ventricle** | 2 | Ventricles vs. everything else | |
|
|
| ### Model Architectures |
|
|
| - **U-Net**: Classic encoder-decoder with skip connections |
| - **U-Net++**: Nested skip pathways for improved feature propagation |
| - **UNETR**: Vision Transformer encoder with CNN decoder |
| - **Swin UNETR**: Hierarchical shifted-window attention |
|
|
| All models trained on **256×256 axial FLAIR images** from 64 patients (Fold 4 training set). |
|
|
| --- |
|
|
| ## 📊 Performance (Fold 4 Validation Results) |
|
|
| ### Multi-Class Segmentation (Dice Score) |
|
|
| | Model | Ventricles | Normal WMH | Abnormal WMH | Mean | |
| |-------|:----------:|:----------:|:------------:|:----:| |
| | **U-Net** | **0.8967** | **0.5935** | **0.6709** | **0.7204** | |
| | U-Net++ | 0.8904 | 0.5881 | 0.6512 | 0.7099 | |
| | UNETR | 0.8401 | 0.4692 | 0.6632 | 0.6575 | |
| | Swin UNETR | 0.8608 | 0.5203 | 0.5920 | 0.6577 | |
|
|
| ### Binary Lesion Segmentation |
|
|
| | Model | Dice | IoU | HD95 (mm) | |
| |-------|:----:|:---:|:---------:| |
| | **U-Net** | **0.7407** | 0.5882 | 32.64 | |
| | U-Net++ | 0.5930 | 0.4215 | 35.12 | |
| | UNETR | 0.6632 | 0.4963 | 40.85 | |
| | Swin UNETR | 0.5841 | 0.4127 | 38.19 | |
|
|
| ### Binary Ventricle Segmentation |
|
|
| | Model | Dice | IoU | HD95 (mm) | |
| |-------|:----:|:---:|:---------:| |
| | **U-Net** | **0.8967** | 0.8130 | 9.52 | |
| | U-Net++ | 0.8904 | 0.8026 | 10.18 | |
| | Swin UNETR | 0.8608 | 0.7560 | 12.73 | |
| | UNETR | 0.8401 | 0.7240 | 14.92 | |
|
|
| *Results are from validation set of Fold 4. See [paper](https://doi.org/10.6084/m9.figshare.30393475) for complete 5-fold statistics.* |
|
|
| --- |
|
|
| ## 🚀 Quick Start |
|
|
| ### Installation |
|
|
| ```bash |
| pip install tensorflow>=2.10.0 nibabel numpy |
| ``` |
|
|
| ### Load and Use Models |
|
|
| ```python |
| from tensorflow import keras |
| from huggingface_hub import hf_hub_download |
| import numpy as np |
| |
| # Download model |
| model_path = hf_hub_download( |
| repo_id="Bawil/MS3SEG", |
| filename="models/multi_class/U-Net_fold4.h5" |
| ) |
| |
| # Load model |
| model = keras.models.load_model(model_path, compile=False) |
| |
| # Prepare your data (256x256 FLAIR image) |
| # image shape: (batch, 256, 256, 1) |
| predictions = model.predict(image) |
| |
| # For multi-class: get class labels |
| pred_classes = np.argmax(predictions, axis=-1) |
| # Classes: 0=background, 1=ventricles, 2=normal WMH, 3=abnormal WMH |
| |
| # For binary: apply threshold |
| pred_binary = (predictions > 0.5).astype(np.uint8) |
| ``` |
|
|
| ### Download All Models for One Scenario |
|
|
| ```python |
| from huggingface_hub import snapshot_download |
| |
| # Download entire scenario folder |
| snapshot_download( |
| repo_id="Bawil/MS3SEG", |
| allow_patterns="models/multi_class/*", |
| local_dir="./ms3seg_models" |
| ) |
| ``` |
|
|
| --- |
|
|
| ## 📝 Input Requirements |
|
|
| - **Format**: NIfTI (.nii.gz) or NumPy array |
| - **Modality**: T2-FLAIR (axial plane) |
| - **Dimensions**: 256 × 256 pixels |
| - **Channels**: 1 (grayscale) |
| - **Preprocessing**: |
| - Co-registered to FLAIR space |
| - Brain-extracted |
| - Intensity normalized to [0, 1] |
| - Voxel spacing: ~0.9 × 0.9 × 5.7 mm³ |
|
|
| See [preprocessing scripts](https://github.com/Mahdi-Bashiri/MS3SEG/tree/main/preprocessing) in our GitHub repository. |
|
|
| --- |
|
|
| ## 📖 Dataset Information |
|
|
| **MS3SEG** is a Multiple Sclerosis MRI dataset with unique **tri-mask annotations**: |
|
|
| - **100 patients** from Iranian cohort (1.5T Toshiba scanner) |
| - **~2000 annotated slices** with expert consensus |
| - **4 annotation classes**: Background, Ventricles, Normal WMH, Abnormal WMH |
| - **Multiple sequences**: T1w, T2w, T2-FLAIR (axial + sagittal) |
|
|
| **Dataset Access:** [Figshare Repository](https://doi.org/10.6084/m9.figshare.30393475) (CC-BY-4.0 License) |
|
|
| --- |
|
|
| ## 🔧 Model Training Details |
|
|
| All models were trained with: |
|
|
| - **Loss Function**: Unified Focal Loss (combining Dice and Focal components) |
| - **Optimizer**: Adam (lr=1e-4) |
| - **Batch Size**: 4 |
| - **Epochs**: 100 (with early stopping, patience=10) |
| - **Data Split**: 64 train / 16 validation patients (Fold 4) |
| - **Framework**: TensorFlow 2.10+ |
|
|
| Complete training configuration available in `config.json`. |
|
|
| --- |
|
|
| ## 📚 Citation |
|
|
| If you use these models in your research, please cite our paper: |
|
|
| ```bibtex |
| @article{bashiri2026ms3seg, |
| title={A Multiple Sclerosis MRI Dataset with Tri-Mask Annotations for Lesion Segmentation}, |
| author={Bashiri Bawil, Mahdi and Shamsi, Mousa and Ghalehasadi, Aydin and Jafargholkhanloo, Ali Fahmi and Shakeri Bavil, Abolhassan}, |
| journal={Scientific Data}, |
| year={2026}, |
| doi={10.6084/m9.figshare.30393475}, |
| publisher={Nature Publishing Group} |
| } |
| ``` |
|
|
| --- |
|
|
| ## 🔗 Resources |
|
|
| - **📄 Paper**: [Scientific Data](https://doi.org/10.6084/m9.figshare.30393475) |
| - **💾 Dataset**: [Figshare](https://doi.org/10.6084/m9.figshare.30393475) |
| - **💻 Code**: [GitHub](https://github.com/Mahdi-Bashiri/MS3SEG) |
| - **📧 Contact**: mehdi.bashiri.b@gmail.com |
|
|
| --- |
|
|
| ## ⚠️ Important Notes |
|
|
| 1. **Fold 4 Only**: These models represent one fold (Fold 4) from our 5-fold cross-validation. They demonstrate representative performance but should not be considered the final "best" models across all folds. |
|
|
| 2. **Research Use**: These models are provided for research purposes. Clinical validation is required before any diagnostic application. |
|
|
| 3. **Data Compatibility**: Models expect preprocessed data matching our pipeline. See [preprocessing documentation](https://github.com/Mahdi-Bashiri/MS3SEG/tree/main/preprocessing). |
|
|
| 4. **Complete Results**: For all 5 folds and comprehensive evaluation, see our [GitHub repository](https://github.com/Mahdi-Bashiri/MS3SEG) and [paper](https://doi.org/10.6084/m9.figshare.30393475). |
|
|
| 5. **Storage Considerations**: Full 5-fold model collection (38GB) is available upon request. These representative Fold 4 models (6GB) are sufficient for most use cases. |
|
|
| --- |
|
|
| ## 📜 License |
|
|
| **Models**: CC-BY-4.0 (same as dataset) |
| **Code**: MIT License (see [GitHub](https://github.com/Mahdi-Bashiri/MS3SEG)) |
|
|
| You are free to use, modify, and distribute these models with appropriate attribution. |
|
|
| --- |
|
|
| ## 🙏 Acknowledgments |
|
|
| Data acquired at Golgasht Medical Imaging Center, Tabriz, Iran. Ethics approval: Tabriz University of Medical Sciences (IR.TBZMED.REC.1402.902). |
|
|
| --- |
|
|
| <div align="center"> |
|
|
| **Made by the MS3SEG Team** |
|
|
| [GitHub](https://github.com/Mahdi-Bashiri/MS3SEG) • [Paper](https://doi.org/10.6084/m9.figshare.30393475) • [Dataset](https://doi.org/10.6084/m9.figshare.30393475) |
|
|
| </div> |
|
|
