Pre-trained models for white matter hyperintensity (WMH) segmentation with explicit distinction between normal periventricular changes and pathological lesions.

Model Description

This repository contains 8 pre-trained deep learning models (4 architectures × 2 training scenarios) for automated WMH segmentation from FLAIR MRI images. The models implement a novel three-class approach that distinguishes between:

Class 0: Background
Class 1: Normal WMH (aging-related periventricular changes)
Class 2: Abnormal WMH (pathologically significant lesions)

This approach addresses the critical challenge of false positive detection in periventricular regions, achieving up to 27.1% improvement in Dice coefficient compared to traditional binary segmentation.

Model Architectures

Architecture	Parameters	Best Dice (3-Class)	Binary Baseline	Improvement
U-Net ⭐	31.0M	0.768	0.497	+54.5%
Attention U-Net	34.9M	0.740	0.486	+52.1%
TransUNet	105.3M	0.700	0.510	+37.3%
DeepLabV3Plus	40.3M	0.586	0.374	+56.7%

⭐ Recommended: U-Net with Scenario 2 (three-class) for optimal performance

Repository Structure

models/
├── unet/models/
│   ├── scenario1_binary_model.h5          # Binary: Background vs Abnormal
│   └── scenario2_multiclass_model.h5      # 3-Class: Background, Normal, Abnormal
├── attention_unet/models/
│   ├── scenario1_binary_model.h5
│   └── scenario2_multiclass_model.h5
├── deeplabv3plus/models/
│   ├── scenario1_binary_model.h5
│   └── scenario2_multiclass_model.h5
└── transunet/models/
    ├── scenario1_binary_model.h5
    └── scenario2_multiclass_model.h5

Quick Start

Installation

pip install huggingface_hub tensorflow numpy nibabel

Download Models

python

from huggingface_hub import hf_hub_download

# Download best performing model (U-Net Three-Class)
model_path = hf_hub_download(
    repo_id="Bawil/wmh_leverage_normal_abnormal_segmentation",
    filename="unet/models/scenario2_multiclass_model.h5"
)

# Load model
from tensorflow.keras.models import load_model
model = load_model(model_path)

Inference Example

python

import numpy as np
from tensorflow.keras.models import load_model

# Load pre-trained model
model = load_model(model_path)

# Prepare input (256x256 grayscale FLAIR MRI, normalized)
# input_image shape: (batch_size, 256, 256, 1)
input_image = preprocess_flair(your_flair_image)

# Run inference
predictions = model.predict(input_image)

# Get class predictions
predicted_classes = np.argmax(predictions, axis=-1)
# 0: Background
# 1: Normal WMH (periventricular)
# 2: Abnormal WMH (pathological)

# Extract pathological lesions only
abnormal_mask = (predicted_classes == 2).astype(np.uint8)

Training Data

Dataset Composition

Local Dataset: 100 MS patients (2,000 FLAIR MRI slices)
- Demographics: 26 males, 74 females
- Age range: 18-68 years
- Scanner: 1.5-Tesla TOSHIBA Vantage
Public Dataset: MSSEG2016 (15 patients, 750 FLAIR slices)

Annotations

Expert annotations by board-certified neuroradiologists (20+ years experience)
Three-class labeling: Background, Normal WMH, Abnormal WMH
Approved by Ethics Committee (IR.TBZMED.REC.1402.902)

Data Split

Training: 80% patients (local) + 60% patients (public)
Validation: 10% patients (local) + 20% patients (public)
Testing: 10% patients (local) + 20% patients (public)
Strategy: Patient-level stratified split (no slice-level leakage)

Model Training

Configuration

Framework: TensorFlow 2.11, Keras
Optimizer: Adam (learning rate: 0.0001)
Loss Functions:
- Scenario 1: Weighted binary cross-entropy
- Scenario 2: Weighted categorical cross-entropy
Epochs: 50 (with early stopping)
Batch Size: 8
Input Size: 256×256×1
Data Augmentation: Rotation, flipping, elastic deformation

Hardware

GPU: NVIDIA RTX 3060 (12GB VRAM)
Training Time: 2-3 hours per model
Inference Time: ~35-40ms per image

Model Performance

Dice Coefficient (Primary Metric)

Model	Scenario 1	Scenario 2	Δ Improvement	p-value	Cohen's d
U-Net	0.497±0.145	0.768±0.124	+0.271	<0.0001	0.564
Attention U-Net	0.486±0.157	0.740±0.133	+0.253	<0.0001	0.442
TransUNet	0.510±0.116	0.700±0.097	+0.190	<0.0001	0.478
DeepLabV3Plus	0.374±0.110	0.586±0.092	+0.212	<0.0001	0.565

Additional Metrics

Hausdorff Distance: 27.4mm (U-Net 3-class) vs 29.8mm (binary)
Precision: Significant improvement in pathological lesion detection
False Positive Reduction: Marked decrease in periventricular regions
Clinical Feasibility: 1.5s total processing time per case (40 slices)

Statistical Validation

Paired t-tests confirm significant improvements (all p < 0.0001)
Effect sizes range from medium (0.44) to large (0.56)
95% confidence intervals reported for all metrics
Wilcoxon signed-rank test for non-parametric validation

Use Cases

Clinical Applications

MS Lesion Quantification: Accurate measurement of disease burden
Differential Diagnosis: Distinguish pathological from normal aging
Longitudinal Monitoring: Track disease progression over time
Treatment Response: Evaluate therapeutic efficacy
Radiological Reporting: Reduce false positive alerts

Research Applications

Baseline Comparisons: Standardized evaluation framework
Method Development: Foundation for advanced segmentation approaches
Multi-center Studies: Protocol for broader validation
Reproducible Research: Complete implementation available

Limitations

Single Modality: Trained on FLAIR MRI only
Scanner Specificity: Primarily 1.5T TOSHIBA data
Disease Focus: Optimized for MS patients
2D Segmentation: Slice-by-slice processing (no 3D context)
Resolution: Fixed 256×256 input size

Model Card

Intended Use

Primary: Automated WMH segmentation for research and clinical decision support
Users: Radiologists, neurologists, researchers, AI developers
Out-of-scope: Not FDA/CE approved; not for standalone clinical diagnosis

Ethical Considerations

Privacy: All data anonymized per HIPAA/GDPR standards
Bias: Limited scanner/protocol diversity may affect generalization
Clinical Validation: Requires expert review before clinical use
Transparency: Complete methodology and code openly available

Model Card Authors

Mahdi Bashiri Bawil, Mousa Shamsi, Ali Fahmi Jafargholkhanloo, Abolhassan Shakeri Bavil

Citation

bibtex

@article{bawil2025wmh,
  title={Incorporating Normal Periventricular Changes for Enhanced Pathological 
         White Matter Hyperintensity Segmentation: On Multi-Class Deep Learning Approaches},
  author={Bawil, Mahdi Bashiri and Shamsi, Mousa and Jafargholkhanloo, Ali Fahmi and 
          Bavil, Abolhassan Shakeri},
  year={2025},
  note={Models: https://huggingface.co/Bawil/wmh_leverage_normal_abnormal_segmentation}
}

License

MIT License - See LICENSE

Additional Resources

📄 Paper: [Under Review]
💻 GitHub Repository: Mahdi-Bashiri/wmh-normal-abnormal-segmentation
📧 Contact: m_bashiri99@sut.ac.ir
🏥 Institution: Sahand University of Technology & Tabriz University of Medical Sciences

Acknowledgments

Golgasht Medical Imaging Center, Tabriz, Iran for providing clinical data
Expert neuroradiologists for manual annotations
Ethics Committee approval: IR.TBZMED.REC.1402.902

Keywords: white matter hyperintensities, FLAIR MRI, medical imaging, deep learning, image segmentation, multiple sclerosis, U-Net, attention mechanisms, transformers, clinical AI