Brats2020 Attention Unet

BraTS2020 Brain Tumor Segmentation: Attention U-Net

This repository contains the champion Attention U-Net model trained from scratch on the BraTS 2020 (Brain Tumor Segmentation Challenge) dataset. The model achieves publication-grade volumetric accuracy and segment boundaries on multimodal brain MRI scans.

• Developer: Bilge
• Model Architecture: Attention U-Net (ResNet-like Convolutional Blocks + Self-Attention Gates)
• Framework: PyTorch (Stable Mixed-Precision training)
• Target Task: Multimodal Brain Tumor Segmentation (FLAIR, T1, T1ce, T2)
• Classes:
  • NCR/NET (Necrotic & Non-Enhancing Tumor Core)
  • Edema (Peritumoral Edema)
  • ET (Enhancing Active Tumor)

---

🏆 Final Performance Metrics

Trained to convergence for 42 epochs with stable log-space Focal-Dice loss, AMP mixed-precision, and Batch Normalization calibration, the model achieves outstanding volumetric validation scores:

Metric	Score (Dice Coefficient)	Clinical Significance
BraTS Mean Score	84.13%	Superior overall tumor delineation.
Whole Tumor (WT)	86.20%	Precise mapping of fluid & vasogenic edema boundaries.
Tumor Core (TC)	80.90%	Excellent localization of the inner tumor core.
Enhancing Tumor (ET)	82.50%	Outstanding detection of highly active ring-enhancing margins.

---

🖼️ Sample Segmentation Visual Output

Below is a sample high-resolution prediction on an unseen validation slice compared side-by-side with the expert Ground Truth annotations:

Attention U-Net Brain MRI Segmentation Comparison

Notice the incredibly clean boundary alignment between the Ground Truth and our Attention U-Net prediction, demonstrating the model's excellent generalization capacity.

---

🧠 Architectural Highlights: The Power of Attention Gates

Unlike a vanilla U-Net, the Attention U-Net integrates Attention Gates (AGs) in the skip connections.

Skip Connection (x) ────► [ Attention Gate ] ────► Concatenate & Decode
                              ▲
Decoder Signal (g) ───────────┘

The gate utilizes the low-resolution coarse feature map from the decoder ($g$) to filter and scale the high-resolution skip connection features ($x$) coming from the encoder. This mechanism forces the model to focus its receptive field on the salient, clinically relevant tumor boundaries while suppressing activations in background brain tissue and noise.

---

🚀 How to Load and Predict in PyTorch

To use this model in PyTorch, ensure you follow the BatchNorm calibration rule during inference to bypass running statistics drift:

import torch
import torch.nn as nn
from src.models import build_model

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 1. Rebuild the model structure
model = build_model(
    model_name="attention_unet",
    in_channels=4,
    out_channels=3,
    encoder_weights="None"
).to(device)

# 2. Load the downloaded weights
ckpt = torch.load("best_attention_final.pt", map_location=device)
model.load_state_dict(ckpt["model_state_dict"])

# 3. Enter Eval Mode but calibrate BatchNorms
model.eval()
model.apply(lambda m: m.train() if isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)) else None)

# Now model is ready for perfect %84.13 volumetric inference!

---

📂 Dataset Information

The model was trained on pre-extracted 2D slices of the BraTS 2020 training dataset. Each slice is a 4-channel tensor of shape (4, 240, 240) containing:

1. FLAIR (Fluid-Attenuated Inversion Recovery)
2. T1 (T1-weighted)
3. T1ce (T1-weighted Contrast-Enhanced)
4. T2 (T2-weighted)

All MRI modalities were normalized using localized foreground z-score normalization prior to training.