brain-mri-segmentation

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.

• Architecture: SegformerForSemanticSegmentation (MiT-B2 encoder, all-MLP decoder)
• Task: binary semantic segmentation, id2label = {0: background, 1: tumor}
• Input: RGB brain-MRI slice; Output: per-pixel class logits at the input resolution
• Base model: nvidia/segformer-b2-finetuned-ade-512-512

Metrics

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.

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.

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

import torch
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation

repo = "kiselyovd/brain-mri-segmentation"
processor = AutoImageProcessor.from_pretrained(repo)
model = AutoModelForSemanticSegmentation.from_pretrained(repo)
model.eval()

image = Image.open("brain_slice.png").convert("RGB")
inputs = processor(images=image, return_tensors="pt")

with torch.no_grad():
    logits = model(**inputs).logits  # (batch, num_labels, H/4, W/4)

# Upsample logits to the original slice size, then take the argmax.
upsampled = torch.nn.functional.interpolate(
    logits,
    size=image.size[::-1],  # (height, width)
    mode="bilinear",
    align_corners=False,
)
mask = 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.

License

MIT - see the LICENSE in the source repository.