ViT for brain tumour MRI classification

google/vit-base-patch16-224 fine-tuned to classify a brain MRI scan as glioma, meningioma, no tumour or pituitary tumour. Made for the MAAI901 Deep Learning for Visual Data course (CNN2ViT workshop). Demo: a Gradio web app in the project's Colab notebook.

Educational project only - not for medical diagnosis.

Training

• Dataset: Kaynaaf/Brain-Tumour-MRI (4,855 training, 857 validation, 1,311 test scans)
• 5 epochs, batch size 32, learning rate 5e-5 with 10% warm-up, mixed precision, light augmentation (flip + tilt up to 10 degrees)
• Kept the epoch with the best validation accuracy

Results on the test set

• All 1,311 test scans: 98.9% accuracy
• The 864 test scans without a near-identical copy in the training data: 98.4% accuracy

Limitations

• 447 of the 1,311 test scans have a practically identical copy in the training data, which inflates the headline score.
• The classes come from different sources: a simple model that only sees image size and colour format (no brain pixels) finds 99% of the "no tumour" scans, so the model may partly rely on such shortcuts.
• Attention maps show it sometimes focuses on the skull or the eyes instead of the tumour.
• Not tested on scans from other hospitals or scanners.