Brats2020 Attention Unet

🏆 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:

🖼️ 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.