Koushim/vit-brain-mri-classifier

🧠 Brain Tumor Classification Using Vision Transformer (ViT)

This repository contains a fine-tuned Vision Transformer (ViT) model trained on a large collection of MRI scans for brain tumor classification. The model classifies MRI images into one of three categories:

• —Glioma
• —Meningioma
• —Tumor (General)

The dataset used includes over 75,000 color-enhanced MRI images, making this model highly capable for research and educational applications in brain tumor detection.

📊 Dataset Information

• —Original Dataset Name: Brain Cancer - MRI dataset
• —Author: Rahman, Md Mizanur (2024)
• —Hosted on: Mendeley Data
• —DOI: 10.17632/mk56jw9rns.1
• —Kaggle Rehost (Colorized): Shuvo Kumar Basak on Kaggle

Note: This dataset is publicly available for non-commercial research use. The model does not include the dataset itself.

🧠 Model Architecture

• —Model Type: Vision Transformer (ViT-B/16)
• —Framework: PyTorch + timm
• —Input Shape: 224x224 RGB
• —Number of Classes: 3
• —Loss Function: CrossEntropyLoss
• —Optimizer: AdamW

🏁 Training Pipeline Summary

1. 1.Image Preprocessing:
2. 2.Resize to 224x224
3. 3.Normalization using ImageNet stats
4. 4.Augmentations: Horizontal/Vertical Flip, ShiftScaleRotate, BrightnessContrast, etc.

1. 1.DataLoader:
2. 2.Stratified Split (Train/Val/Test)
3. 3.PyTorch Dataset and DataLoader classes

1. 1.Model:
2. 2.Loaded ViT using timm.create_model('vit_base_patch16_224', pretrained=True)
3. 3.Modified the classifier head to match 3 output classes

1. 1.Training:
2. 2.Trained using mixed precision (torch.cuda.amp)
3. 3.Tracked using tqdm

1. 1.Saving:
2. 2.Model saved as pytorch_model.bin
3. 3.Configuration saved as config.json

🔍 Intended Use

This model is designed for:

• —Educational purposes (deep learning and medical imaging)
• —Research in brain tumor classification using transformers
• —Demonstrating the power of ViT on colorized medical datasets

⚠️ Not intended for clinical use or deployment without regulatory approval and further validation.

🚀 Inference Example (Python)

from timm import create_model
import torch
from torchvision import transforms
from PIL import Image

# Load model
model = create_model('vit_base_patch16_224', pretrained=False, num_classes=3)
model.load_state_dict(torch.load("pytorch_model.bin"))
model.eval()

# Transform
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5]*3, std=[0.5]*3),
])

# Inference
image = Image.open("example_mri.jpg").convert("RGB")
tensor = transform(image).unsqueeze(0)
output = model(tensor)
pred = torch.argmax(output, dim=1)
print("Predicted class:", pred.item())