Binary brain-tumor segmentation on MRI slices, built on a fine-tuned SegFormer-B2 semantic-segmentation transformer. Given a single FLAIR brain-MRI slice, the model predicts a per-pixel mask separating tumor tissue (1) from background (0). It was trained on the Mateusz Buda LGG (TCGA) dataset with a strict patient-level split to prevent data leakage, and a hand-rolled U-Net is included in the project as a reproducible baseline for comparison.
Evaluated on the held-out test split (387 slices) from 11 patients never seen during training. The main SegFormer-B2 model is reported alongside the U-Net baseline trained on the same split.
Model
Dice
IoU
Pixel accuracy
SegFormer-B2 (main)
65.5%
66.2%
99.73%
U-Net (baseline)
51.9%
57.7%
99.66%
Dice and IoU are macro-averaged over the test set; the metric values are reproduced verbatim from reports/metrics.json and reports/metrics_baseline.json in the source repository.
Visualizations
Qualitative results
Real predictions from this model on tumor-bearing test slices. Each row shows the input MRI slice, the ground-truth tumor mask (red), and the model's predicted mask (cyan), with the per-image Dice score in the title. Both mask columns are drawn over a grayscale view of the same slice so the contours stay legible.
Qualitative segmentation results
Per-image Dice on tumor-bearing slices is high (roughly 0.88-0.96), since these are clear, well-defined lesions. The dataset-wide mean of 0.655 is lower because it also averages in tumor-free slices, where a single false-positive pixel drives Dice toward zero, and harder, low-contrast cases.
Metrics comparison
SegFormer-B2 versus the U-Net baseline across Dice, IoU, and pixel accuracy. The transformer wins decisively on the region-overlap metrics (Dice / IoU); pixel accuracy is near-saturated for both because tumor pixels are a small fraction of each slice.
Metrics comparison
Medical disclaimer
This model is provided for research and educational purposes only. It is not a medical device and has not been validated for clinical use. Its predictions must not be used for diagnosis, treatment, or any clinical decision-making. Always consult a qualified medical professional.
Usage
python
1import torch
2from PIL import Image
3from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation
45repo ="kiselyovd/brain-mri-segmentation"6processor = AutoImageProcessor.from_pretrained(repo)7model = AutoModelForSemanticSegmentation.from_pretrained(repo)8model.eval()910image = Image.open("brain_slice.png").convert("RGB")11inputs = processor(images=image, return_tensors="pt")1213with torch.no_grad():14 logits = model(**inputs).logits # (batch, num_labels, H/4, W/4)1516# Upsample logits to the original slice size, then take the argmax.17upsampled = torch.nn.functional.interpolate(18 logits,19 size=image.size[::-1],# (height, width)20 mode="bilinear",21 align_corners=False,22)23mask = upsampled.argmax(dim=1)[0]# 2-D tensor; 1 = tumor, 0 = background
mask is a 2-D tensor aligned to the input slice, with 1 marking predicted tumor pixels and 0 background.
Training data
Trained on LGG MRI Segmentation (TCGA) - 110 patients and 3,929 paired FLAIR slices with binary tumor masks from The Cancer Genome Atlas. The source repository performs a patient-level 80/10/10 split (3,133 train / 409 val / 387 test) so no patient appears in more than one partition.
Source code
GitHub Repository - full training, evaluation, serving, and plotting code, including the scripts/make_plots.py script that produced the visualizations above.
