mahdisabetkish/brainssl-tumour-segmentation

brainssl: bilateral U-Net for brain tumour segmentation

A residual 3D U-Net that segments glioma sub-regions in four-contrast brain MRI (FLAIR, T1, T1 with gadolinium, T2). Each encoder level can read the features of its mirror image across the brain's mid-sagittal plane, and the network was pretrained without labels by reconstructing hidden tissue from its contralateral twin, then fine-tuned on annotated scans.

Code, training pipeline and the full evaluation: https://github.com/mahdisabetkish/brainssl. Interactive results: https://mahdisabetkish.github.io/brainssl/.

Files

File	What it is
segmentation.onnx	one sliding window: input image (N, 4, 160, 128, 128), output logits (N, 3, 160, 128, 128)
segmentation.json	everything around the window: crop, normalisation, window placement, blending, thresholds
model.safetensors, config.json	PyTorch weights for brainssl.smri.model.UNet3D

The three output channels are whole tumour, tumour core and enhancing tumour, each through its own sigmoid; the label map is 1 oedema, 2 non-enhancing core, 3 enhancing tumour.

Use

Without Python, with the Rust binary from the repository:

brainssl-infer --ort-lib /path/to/libonnxruntime.so segment \
    --model . --input scan.nii.gz --output scan_seg.nii.gz

From Python, the brainssl package reproduces the same pipeline (crop about the midline, z-score per contrast, sliding windows with Gaussian blending, mirror averaging).

Training data

Medical Segmentation Decathlon Task01 (BraTS 2016 and 2017 scans, CC BY-SA 4.0): pretraining on 576 scans (310 annotated scans with their labels ignored, plus the 266 released without labels), fine-tuning on the 310 annotated training scans of a fixed split. Validation and test scans were never used for training.

Evaluation

Held-out test split of 96 annotated scans, Dice (higher is better):

Training labels	Model	WT	TC	ET
100% (310 scans)	this model	0.899	0.853	0.840
10% (31 scans)	same method	0.885	0.802	0.802
10% (31 scans)	U-Net from scratch	0.872	0.799	0.790

Limitations

Trained on skull-stripped scans registered to the SRI24 template at 1 mm, as released by the BraTS organisers; other preprocessing will need the same steps first. The mid-sagittal plane is estimated per scan, but strong midline shift from a large tumour weakens the left-right correspondence the model relies on. This is a research model and not a medical device.