classicmaster/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

Metric	Score
Test Loss	0.1926
Dice (Tumor Classes)	0.7293
IoU (Tumor Classes)	0.6164
Voxel Accuracy	0.9900

Per-Class Dice

Class	Dice Score
Necrotic / Non-Enhancing Core	0.7503
Edema	0.8046
Enhancing Tumor	0.8414

Composite Region Dice (standard BraTS regions)

Region	Dice Score
Whole Tumor (WT)	0.9181
Tumor Core (TC)	0.8678
Enhancing Tumor (ET)	0.8414

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

Files in This Repository

File	Description
model.py	Standalone AttentionUNet3D architecture definition (required to load the weights)
best_model.pth	Trained model weights (raw, full precision)
best_model_ema.pth	EMA-averaged weights (typically more stable for inference)
model.safetensors	Weights in safetensors format
inference.py	Standalone script to run inference on a preprocessed .npy MRI volume
evaluation_report.md	Full evaluation report generated during training
final_summary.csv	Test metrics as CSV
report_figures/	Training curves, prediction vs. ground-truth comparisons, overlay visualizations
inference_output/	Sample inference result on an unseen test volume

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.