mastiskhnet-brain-tumor-segmentation

MastiskhNet — 3D Attention U-Net for Brain Tumor Segmentation

MastiskhNet is a 3D Attention U-Net trained on the BraTS 2020 dataset to segment brain tumors from multi-modal MRI volumes (FLAIR, T1, T1ce, T2). It outputs a voxel-wise segmentation into background, necrotic/non-enhancing core, edema, and enhancing tumor.

This repo contains the trained weights, standalone architecture code, and an inference script, built as part of the MastiskhNet project (backend, doctor portal, and mobile app for AI-assisted brain tumor analysis).

Model Details

• Architecture: 3D Attention U-Net (4-level encoder/decoder, GroupNorm + SiLU, additive attention gates on every skip connection)
• Input: 4-channel MRI volume (FLAIR, T1, T1ce, T2), resized to 128×128×128, Z-score normalized
• Output: 4-class voxel segmentation (Background, Necrotic/Core, Edema, Enhancing Tumor)
• Parameters: ~22.7M
• Framework: PyTorch
• Training data: BraTS 2020 (369 patients, 70/15/15 train/val/test split)
• Loss: Combined Dice + Focal Loss
• Training setup: AdamW optimizer, Cosine Annealing Warm Restarts scheduler, Automatic Mixed Precision (AMP), Exponential Moving Average (EMA) of weights, gradient clipping, 100 epochs

Performance

Evaluated on the held-out BraTS 2020 test split (56 batches), using EMA weights + test-time augmentation.

Overall Test Metrics

Per-Class Dice

Composite Region Dice (standard BraTS regions)

Training converged over 100 epochs, reaching a best validation Dice of 0.7139 (validation loss 0.1932).

Files in This Repository

Usage

from model import AttentionUNet3D
import torch

model = AttentionUNet3D()
checkpoint = torch.load("best_model.pth", map_location="cpu")

# Checkpoints saved via the training loop are wrapped in a dict
state_dict = checkpoint["model_state_dict"] if "model_state_dict" in checkpoint else checkpoint
model.load_state_dict(state_dict, strict=False)
model.eval()

# volume: torch.Tensor of shape (1, 4, 128, 128, 128) — FLAIR, T1, T1ce, T2, Z-score normalized
with torch.no_grad():
    logits = model(volume)
    prediction = torch.argmax(logits, dim=1)  # (1, 128, 128, 128), values 0-3

Or use the provided script:

python inference.py --volume path/to/volume.npy --model best_model.pth

Loading from safetensors

from model import AttentionUNet3D
from safetensors.torch import load_file

model = AttentionUNet3D()
model.load_state_dict(load_file("model.safetensors"))
model.eval()

Preprocessing Pipeline

Volumes must be preprocessed identically to training before inference:

1. Load the four MRI modalities (FLAIR, T1, T1ce, T2) and segmentation mask via nibabel
2. Crop to the brain bounding box
3. Resize every modality to 128×128×128 (trilinear for MRI, nearest-neighbor for masks)
4. Z-score normalize each modality (mask excluded)
5. Remap BraTS labels {0, 1, 2, 4} → contiguous {0, 1, 2, 3}
6. Stack the four modalities into a single (4, 128, 128, 128) tensor

Limitations

• Trained exclusively on BraTS 2020 (adult glioma MRI); performance on other tumor types, pediatric scans, or different scanner protocols is unverified.
• Requires all four MRI modalities (FLAIR, T1, T1ce, T2) registered and skull-stripped, matching BraTS preprocessing conventions.
• This model is a research artifact and has not been clinically validated. It is not intended for diagnostic use.

Citation

If you use this model, please cite the BraTS 2020 dataset:

Menze et al., "The Multimodal Brain Tumor Image Segmentation Benchmark (BRATS)", IEEE Transactions on Medical Imaging, 2015.