MohanadKombar/brain-tumor-unetr-brats

🧠 UNETR (Vision Transformer) for Brain Tumor Segmentation (BraTS 2021 + 2024)

This repository contains pre-trained weights, architecture definition, and standalone inference scripts for a UNETR (3D Vision Transformer) network trained for volumetric brain tumor segmentation on multi-modal MRI scans.

• GitHub Repository: https://github.com/MohanadMahran/Brain_Tumor_Benchmark
• Hugging Face Model: MohanadKombar/brain-tumor-unetr-brats

---

📐 Architecture Summary

• Architecture: UNETR (3D Vision Transformer Encoder with CNN Decoder)
• Total Parameters: 19,327,907 (~19.33M)
• Normalization: LayerNorm (Transformer) and GroupNorm (Decoder)
• Transformer Config: Embedding dimension 384, 6 transformer layers, 8 attention heads, $16\times 16\times 16$ patch tokenization
• Input Modalities: 4-channel MRI (T1, T1Gd, T2, FLAIR)
• Output Sub-regions: 3 overlapping sigmoid target channels:
  • Channel 0: Enhancing Tumor (ET) (Label 4)
  • Channel 1: Tumor Core (TC) (Labels 1 + 4)
  • Channel 2: Whole Tumor (WT) (Labels 1 + 2 + 4)

---

🏋️ Training Summary

• Dataset: Composite dataset derived from BraTS 2021 (1,000 train cases) and BraTS 2024 Adult Glioma (560 train cases) after cross-year MD5 deduplication (1,560 total training cases, 391 validation cases).
• Epochs Trained: 500 total epochs (completed without early stopping). Reached peak validation performance at Epoch 469.
• Loss Function: Combined DiceCE Loss (Soft Dice + Cross-Entropy).
• Optimizer & Scheduler: AdamW ($\text{LR}=1.0 \times 10^{-4}$, weight decay $1.0\times 10^{-5}$) with 15-epoch linear warmup and Cosine Annealing decay down to $\text{min_lr}=1.0\times 10^{-6}$.
• Peak Training VRAM: 7,026 MiB (~7.03 GB) on NVIDIA A100-SXM4-40GB (optimized with PyTorch Gradient Checkpointing).

---

📊 Empirical Evaluation & Benchmark Results

Performance Summary

Dataset / Cohort	Mean Dice	ET Dice	TC Dice	WT Dice	Mean HD95 (mm)
BraTS Validation (391 cases)	0.8362	0.7629	0.8277	0.9177	13.95
UPenn-GBM Out-of-Distribution (30 cases)	0.4749 ± 0.3007	0.4454 ± 0.3223	0.4976 ± 0.3387	0.4818 ± 0.2887	87.71 ± 27.26

• Generalization Gap: 36.13% performance drop when evaluated out-of-distribution on UPenn-GBM (0.8362 → 0.4749).
• Inference Time: 1.84s ± 0.19s per case (sliding window inference).
• Evaluation Peak VRAM: 1,979 MB (~1.98 GB).

---

💻 Usage & Inference

Installation

pip install torch monai nibabel huggingface_hub

Quick Start Inference Code

import torch
from inference import load_model, predict

# 1. Load pre-trained model from Hugging Face Hub (or local directory)
model, metadata = load_model("MohanadKombar/brain-tumor-unetr-brats")

# 2. Prepare 4-channel input volume of shape (1, 4, D, H, W)
# Input modalities must be ordered: [T1, T1Gd, T2, FLAIR]
dummy_input = torch.randn(1, 4, 128, 128, 128)

# 3. Run sliding-window inference
probabilities = predict(model, dummy_input) # Output shape: (1, 3, 128, 128, 128)

# 4. Threshold at 0.5 for binary segmentation masks
pred_binary = (probabilities > 0.5).cpu().numpy()[0] # [0: ET, 1: TC, 2: WT]

⚠️ Crucial Data Preprocessing Requirement

Inputs MUST be preprocessed prior to running inference to produce valid segmentations:

1. Percentile Clipping: Clip intensity values to [0.5, 99.5] percentile range per modality.
2. Brain Mask Extraction: Compute non-zero brain mask (intensity > 0 across modalities).
3. Z-Score Normalization: Normalize voxels per channel within the brain mask: $\frac{x - \mu}{\sigma}$.

---

📜 Citation & License

• License: MIT License
• Reference: Fair Benchmarking of 3D CNN vs Vision Transformer Architectures for Volumetric Brain Tumor Segmentation.