MohanadKombar/brain-tumor-unet3d-brats

🧠 3D U-Net for Brain Tumor Segmentation (BraTS 2021 + 2024)

This repository contains pre-trained weights, architecture definition, and standalone inference scripts for a 3D U-Net (3D Convolutional Neural 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-unet3d-brats

📐 Architecture Summary

• —Architecture: 3D U-Net (4-level encoder-decoder with residual skip connections)
• —Total Parameters: 19,992,603 (~19.99M)
• —Normalization: GroupNorm (groups=4) — batch-size invariant
• —Channel Progression: [60, 120, 240, 480]
• —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 419.
• —Loss Function: Combined DiceCE Loss (Soft Dice + Cross-Entropy).
• —Optimizer & Scheduler: AdamW ($\text{LR}=3.0 \times 10^{-4}$, weight decay $1.0\times 10^{-5}$) with 10-epoch linear warmup and Cosine Annealing decay down to $\text{min\_lr}=1.0\times 10^{-6}$.
• —Peak Training VRAM: 19,025 MiB (~19.03 GB) on NVIDIA A100-SXM4-40GB.

📊 Empirical Evaluation & Benchmark Results

Performance Summary

• —Generalization Gap: 35.08% performance drop when evaluated out-of-distribution on UPenn-GBM (0.8694 → 0.5186).
• —Inference Time: 4.58s ± 0.42s per case (sliding window inference).
• —Evaluation Peak VRAM: 3,718 MB (~3.72 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-unet3d-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. 1.Percentile Clipping: Clip intensity values to [0.5, 99.5] percentile range per modality.
2. 2.Brain Mask Extraction: Compute non-zero brain mask (intensity > 0 across modalities).
3. 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.