BraTS Segmentation Using 3D U-Net

📁 Dataset Info

The dataset used is the BraTS 2020 multimodal brain tumor segmentation dataset, available on Kaggle:
https://www.kaggle.com/datasets/awsaf49/brats20-dataset-training-validation

• All scans are in .nii.gz (NIfTI) format.
• Each sample includes 4 MRI modalities:
  • T1: Native
  • T1CE: Post-contrast T1-weighted
  • T2: T2-weighted
  • FLAIR: Fluid Attenuated Inversion Recovery
• Ground truth labels:
  • 0: Background
  • 1: Necrotic/Non-enhancing tumor core (NCR/NET)
  • 2: Peritumoral edema (ED)
  • 4: GD-enhancing tumor (ET)

⚙️ Tools & Libraries Used

• numpy
• nibabel
• glob
• matplotlib
• tensorflow and keras
• sklearn.preprocessing.MinMaxScaler

🧠 Segmentation Pipeline

1. Step 1: Data Preparation
   • Loaded MRI modalities using nibabel
   • Normalized volume values using MinMaxScaler
   • Cropped original shape 240×240×155 to 128×128×128 for uniform training
   • Filtered out samples where tumor covered < 1% of the brain volume
   • Split data into 75% training and 25% validation using splitfolders
2. Step 2: Custom Data Generator
   • Keras' ImageDataGenerator doesn't support .npy files
   • Created a generator that:
     • Loads preprocessed data from disk
     • Applies shuffling and batch generation
     • Returns 3D MRI volume and corresponding mask
3. Step 3: Define the 3D U-Net Model
   • Extended standard 2D U-Net to 3D using:
     • Conv3D, MaxPooling3D, and UpSampling3D
4. Step 4: Training and Prediction
   • Trained using the custom data generator
   • Used Dice coefficient and Categorical Focal Loss
   • Visualized predictions using matplotlib (random 2D slices)

📊 Results : 6GB rtx 3050 - 10 Epochs

Training Loss

Mean IoU: 61.34 (10 Epochs)

Prediction Samples

Training vs Validation Loss

Training vs Validation Accuracy

📌 Future Improvements

• Training the model for more epochs could further improve segmentation accuracy and stability.
• Experimenting with advanced loss functions such as Dice + Focal Loss or Tversky Loss may enhance performance on complex tumor structures.
• Currently, only the training set (~30 GB) was used and split into training and validation. Future work can utilize the full BraTS dataset (including the validation and testing sets) for more robust training and evaluation.

🔗 References

• 3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation (arXiv)
• YouTube Channel: Digital Sreeni
• Real Python: Introduction to Python Generators
• Focal and Dice Loss Combined: Bioinformatics Advances