You can deploy this model and start earning money today!

MohanadKombar/
brain-tumor-unet3d-brats

Name: brain-tumor-unet3d-brats
Availability: InStock

pytorch

unet3d

medical-imaging

brain-tumor-segmentation

Views

No views yet

🧠 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

Dataset / Cohort	Mean Dice	ET Dice	TC Dice	WT Dice	Mean HD95 (mm)
BraTS Validation (391 cases)	0.8694	0.8211	0.8576	0.9218	13.95
UPenn-GBM Out-of-Distribution (30 cases)	0.5186 ± 0.3223	0.5188 ± 0.3809	0.4855 ± 0.3658	0.5514 ± 0.2616	90.58 ± 31.52

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

python

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:

Percentile Clipping: Clip intensity values to [0.5, 99.5] percentile range per modality.
Brain Mask Extraction: Compute non-zero brain mask (intensity > 0 across modalities).
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.