Koushim/vit-brain-mri-classifier

1---2license: mit3tags:4- medical5---6# 🧠 Brain Tumor Classification Using Vision Transformer (ViT)7 8This 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:9 10- **Glioma**11- **Meningioma**12- **Tumor (General)**13 14The dataset used includes over **75,000 color-enhanced MRI images**, making this model highly capable for research and educational applications in brain tumor detection.15 16---17 18## 📊 Dataset Information19 20- **Original Dataset Name**: Brain Cancer - MRI dataset21- **Author**: Rahman, Md Mizanur (2024)22- **Hosted on**: [Mendeley Data](https://data.mendeley.com/datasets/mk56jw9rns/1)  23- **DOI**: [10.17632/mk56jw9rns.1](https://doi.org/10.17632/mk56jw9rns.1)  24- **Kaggle Rehost (Colorized)**: [Shuvo Kumar Basak on Kaggle](https://www.kaggle.com/datasets/shuvokumarbasakbd/brain-cancer-mri-colorized-dataset)25 26> **Note:** This dataset is publicly available for non-commercial research use. The model does not include the dataset itself.27 28---29 30## 🧠 Model Architecture31 32- **Model Type**: Vision Transformer (ViT-B/16)33- **Framework**: PyTorch + [timm](https://github.com/huggingface/pytorch-image-models)34- **Input Shape**: 224x224 RGB35- **Number of Classes**: 336- **Loss Function**: CrossEntropyLoss37- **Optimizer**: AdamW38 39---40 41## 🏁 Training Pipeline Summary42 431. **Image Preprocessing**:44   - Resize to 224x22445   - Normalization using ImageNet stats46   - Augmentations: Horizontal/Vertical Flip, ShiftScaleRotate, BrightnessContrast, etc.47 482. **DataLoader**:49   - Stratified Split (Train/Val/Test)50   - PyTorch `Dataset` and `DataLoader` classes51 523. **Model**:53   - Loaded ViT using `timm.create_model('vit_base_patch16_224', pretrained=True)`54   - Modified the classifier head to match 3 output classes55 564. **Training**:57   - Trained using mixed precision (`torch.cuda.amp`)58   - Tracked using `tqdm` 59 605. **Saving**:61   - Model saved as `pytorch_model.bin`62   - Configuration saved as `config.json`63 64---65 66## 🔍 Intended Use67 68This model is designed for:69 70- Educational purposes (deep learning and medical imaging)71- Research in brain tumor classification using transformers72- Demonstrating the power of ViT on colorized medical datasets73 74⚠️ **Not intended for clinical use** or deployment without regulatory approval and further validation.75 76---77 78## 🚀 Inference Example (Python)79 80```python81from timm import create_model82import torch83from torchvision import transforms84from PIL import Image85 86# Load model87model = create_model('vit_base_patch16_224', pretrained=False, num_classes=3)88model.load_state_dict(torch.load("pytorch_model.bin"))89model.eval()90 91# Transform92transform = transforms.Compose([93    transforms.Resize((224, 224)),94    transforms.ToTensor(),95    transforms.Normalize(mean=[0.5]*3, std=[0.5]*3),96])97 98# Inference99image = Image.open("example_mri.jpg").convert("RGB")100tensor = transform(image).unsqueeze(0)101output = model(tensor)102pred = torch.argmax(output, dim=1)103print("Predicted class:", pred.item())