|
|
--- |
|
|
license: mit |
|
|
tags: |
|
|
- medical |
|
|
--- |
|
|
# 🧠 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](https://data.mendeley.com/datasets/mk56jw9rns/1) |
|
|
- **DOI**: [10.17632/mk56jw9rns.1](https://doi.org/10.17632/mk56jw9rns.1) |
|
|
- **Kaggle Rehost (Colorized)**: [Shuvo Kumar Basak on Kaggle](https://www.kaggle.com/datasets/shuvokumarbasakbd/brain-cancer-mri-colorized-dataset) |
|
|
|
|
|
> **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](https://github.com/huggingface/pytorch-image-models) |
|
|
- **Input Shape**: 224x224 RGB |
|
|
- **Number of Classes**: 3 |
|
|
- **Loss Function**: CrossEntropyLoss |
|
|
- **Optimizer**: AdamW |
|
|
|
|
|
--- |
|
|
|
|
|
## 🏁 Training Pipeline Summary |
|
|
|
|
|
1. **Image Preprocessing**: |
|
|
- Resize to 224x224 |
|
|
- Normalization using ImageNet stats |
|
|
- Augmentations: Horizontal/Vertical Flip, ShiftScaleRotate, BrightnessContrast, etc. |
|
|
|
|
|
2. **DataLoader**: |
|
|
- Stratified Split (Train/Val/Test) |
|
|
- PyTorch `Dataset` and `DataLoader` classes |
|
|
|
|
|
3. **Model**: |
|
|
- Loaded ViT using `timm.create_model('vit_base_patch16_224', pretrained=True)` |
|
|
- Modified the classifier head to match 3 output classes |
|
|
|
|
|
4. **Training**: |
|
|
- Trained using mixed precision (`torch.cuda.amp`) |
|
|
- Tracked using `tqdm` |
|
|
|
|
|
5. **Saving**: |
|
|
- Model saved as `pytorch_model.bin` |
|
|
- 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) |
|
|
|
|
|
```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()) |
