imbenjita/BrainTumorClassifier

Brain Tumor MRI Classification: SegFormer-B5 + ConvNeXt-Base

Live Demo (Space) 🤗

Research demo, not a medical device. Do not use this model to diagnose or to make clinical decisions.

Introduction

This repository contains a two-stage model that classifies a brain MRI slice as glioma, meningioma, pituitary tumor or no tumor. It is trained on the BRISC 2025 T1-weighted dataset.

Architecture.

1. SegFormer-B5 (nvidia/segformer-b5-finetuned-ade-640-640, Mix Transformer encoder + All-MLP decoder) is fine-tuned as a binary tumor segmenter (background / tumor) with a frozen encoder. It outputs a per-pixel tumor probability map.
2. ConvNeXt-Base (facebook/convnext-base-224) is the classifier. Its input is the RGB scan plus the tumor map as a fourth channel. The first three stages are frozen. Features from the last stage are pooled and passed to a linear head.

Minimizing false negatives. Missing a tumor is treated as more costly than a false alarm:

• Training loss: class-weighted cross-entropy plus an asymmetric detection loss. A tumor called "no tumor" costs 5× a false alarm.
• Decision rule: a minimum-expected-cost rule instead of plain argmax. A slice is cleared as no_tumor only when P(no_tumor) is above roughly 0.83. Every other slice gets the most likely tumor type.

Two classifier variants.

Variant	Checkpoint	Classifier input	Needs a mask at inference?
pred_seg (main model)	seg_best.pt + cls_pred_seg_best.pt	scan + SegFormer-predicted tumor map	No
gt_seg (upper bound)	cls_gt_seg_best.pt	scan + expert tumor mask	Yes

The gt_seg model shows how well the classifier can do when the tumor location is perfect. An empty mask is treated as no_tumor.

Data handling.

• Train images that are byte-identical to a test image, or repeated within train, were removed (SHA-256) to prevent leakage.
• The validation set was cut from train in blocks of 20 consecutive indices per class and plane. This is a heuristic, because the dataset has no patient IDs, so slices of the same patient may still appear in both train and validation.
• The test set is the official BRISC test split (1,000 slices).

Figure 1. Left: one example of every class and anatomical plane (red outline = ground-truth mask). Right: one training batch.

Evaluation Results

All numbers are on the official BRISC test split (1,000 slices: 254 glioma, 306 meningioma, 300 pituitary, 140 no tumor; 398 axial, 305 coronal, 297 sagittal). Metrics use the cost-aware decision rule. On this test set the cost-aware rule and plain argmax gave identical predictions.

Metric (test)	pred_seg (scan only)	gt_seg (scan + true mask)
Accuracy	95.9	96.7
Macro F1	96.3	97.1
Tumor sensitivity (tumor vs no tumor)	100.0 (860/860)	100.0 (860/860)
Specificity (no-tumor slices cleared)	98.6 (138/140)	100.0 (140/140)
False negatives (tumor called no tumor)	0	0
False alarms (no tumor called tumor)	2	0

Per-class results (test, cost-aware rule)

Class	pred_seg precision	pred_seg recall	pred_seg F1	gt_seg precision	gt_seg recall	gt_seg F1
glioma	96.7	92.9	94.8	97.1	91.3	94.1
meningioma	94.1	93.1	93.6	92.5	97.1	94.7
pituitary	95.2	100.0	97.6	99.3	99.3	99.3
no_tumor	100.0	98.6	99.3	100.0	100.0	100.0

Accuracy by anatomical plane (test)

Plane	Slices	pred_seg accuracy	gt_seg accuracy
Axial	398	95.7	98.2
Coronal	305	96.4	95.1
Sagittal	297	95.6	96.3

Tumor sensitivity was 100% in every plane for both variants. The most common error is glioma versus meningioma, in both directions.

Confusion matrices

Red box = false negatives (a tumor called no tumor). The right panel collapses the four classes into tumor / no tumor.

Main model (pred_seg), test set

Upper bound (gt_seg), test set

More confusion matrices (argmax rule and validation set)

Interpretability

Grad-CAM (on the last ConvNeXt stage) and SHAP (GradientExplainer on P(tumor)) are provided for both variants. Heat maps show evidence for the predicted class.

Main model (pred_seg): Grad-CAM, random correctly classified test slices (last column = SegFormer tumor map)

Main model (pred_seg): SHAP (red = evidence for tumor, blue = against)

Upper bound (gt_seg) Grad-CAM and SHAP

How to read these. The SHAP panels on the tumor-map and mask channels show that the classifier uses the segmented region. Grad-CAM often also highlights areas outside the tumor, such as the skull edge, other anatomy or the image border, even for correct predictions. These maps show where the model got evidence, not a verified tumor localization. Treat them as a debugging aid, not as proof that the model "looks at the tumor".

Usage

Files needed in one folder: brisc_segformer_convnext.py (model and prediction code), seg_best.pt, cls_pred_seg_best.pt, and optionally cls_gt_seg_best.pt. Only the trainable weights are stored in the .pt files. The frozen base weights are downloaded from the Hub on first load.

import os
os.environ["BRISC_OUTPUT_DIR"] = "."   # folder holding the .pt files

import brisc_segformer_convnext as brisc

# Scan only: SegFormer finds the tumor, ConvNeXt classifies it.
segmenter, classifier = brisc.load_trained_models(brisc.MODE_PREDICTED_MASK)
result = brisc.predict_tumor_type("scan.jpg", segmenter, classifier)
print(result["tumor_type"], result["tumor_probability"])

# Scan + a tumor mask you already have (any non-zero pixel = tumor).
_, true_mask_classifier = brisc.load_trained_models(brisc.MODE_TRUE_MASK)
result = brisc.predict_tumor_type_with_mask("scan.jpg", "mask.png", true_mask_classifier)

The returned dictionary contains:

Key	Meaning
tumor_type	Recommended label from the false-negative-minimizing rule
argmax_type	Plain highest-probability class
tumor_detected	True unless the label is no_tumor
tumor_probability	1 − P(no_tumor)
probabilities_pct	Per-class probabilities in percent
tumor_area_pct	Share of the slice covered by the tumor map or mask
figure	Path to the saved scan / tumor map / Grad-CAM picture

Preprocessing: the slice is padded to a square and resized (512×512 for SegFormer, 320×320 for ConvNeXt). The segmenter must be used at the size it was trained at.

Figure 2. Left: scan-only prediction (pred_seg). Right: prediction with a provided mask (gt_seg).

Per-slice test predictions with class probabilities are available in test_predictions.csv (pred_seg) and gt_test_predictions.csv (gt_seg).

Limitations

• Not for clinical use. The model was trained and tested on one public dataset of 2D T1 slices. It has not been validated on other scanners, sequences, hospitals or populations, and 2D slices carry no 3D context.
• Possible optimistic estimates. BRISC has no patient IDs. Exact duplicates were removed, but slices from the same patient can still appear in both train and test, which would inflate the results above.
• The 100% sensitivity is on 860 tumor slices. It does not guarantee that no tumor will ever be missed. The cost-aware rule trades false negatives for false alarms, so borderline scans will be flagged as tumors.
• Only four labels. Any other pathology, an unusual scan, or a non-brain image will still be forced into one of the four classes.
• Segmentation quality is moderate. The predicted tumor map is only an approximate localization, and the classifier is more accurate when given true masks (gt_seg).
• Grad-CAM can highlight non-tumor regions (see the interpretability note above).

License

The model code in this repository and the trained checkpoints are released for research use. Please check the licenses of what they build on before any other use:

• SegFormer (nvidia/segformer-b5-finetuned-ade-640-640): see the model card for its license, which may restrict commercial use.
• ConvNeXt (facebook/convnext-base-224): see the model card for its license.
• BRISC 2025 dataset: see the dataset's terms.

Citation

If you use this model, please also cite the dataset and the backbones:

@article{fateh2025brisc,
  title={Brisc: Annotated dataset for brain tumor segmentation and classification with swin-hafnet},
  author={Fateh, Amirreza and Rezvani, Yasin and Moayedi, Sara and Rezvani, Sadjad and Fateh, Fatemeh and Fateh, Mansoor and Abolghasemi, Vahid},
  journal={arXiv preprint arXiv:2506.14318},
  year={2025}
}

@inproceedings{xie2021segformer,
  title={SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers},
  author={Xie, Enze and Wang, Wenhai and Yu, Zhiding and Anandkumar, Anima and Alvarez, Jose M and Luo, Ping},
  booktitle={NeurIPS},
  year={2021}
}

@inproceedings{liu2022convnet,
  title={A ConvNet for the 2020s},
  author={Liu, Zhuang and Mao, Hanzi and Wu, Chao-Yuan and Feichtenhofer, Christoph and Darrell, Trevor and Xie, Saining},
  booktitle={CVPR},
  year={2022}
}

Contact

Questions and issues: please open a discussion on this model's Community tab.