|
|
--- |
|
|
license: mit |
|
|
tags: |
|
|
- medical-imaging |
|
|
- mri |
|
|
- brain-tumor |
|
|
- pytorch |
|
|
- mojo |
|
|
- image-classification |
|
|
datasets: |
|
|
- kaggle-brain-tumor-mri |
|
|
- ixi |
|
|
metrics: |
|
|
- accuracy |
|
|
- f1 |
|
|
- precision |
|
|
- recall |
|
|
pipeline_tag: image-classification |
|
|
library_name: pytorch |
|
|
--- |
|
|
|
|
|
# MRI Brain Tumor Classification Models |
|
|
|
|
|
[](https://github.com/Meidverse/COMMRI) |
|
|
[](LICENSE) |
|
|
|
|
|
This repository contains trained deep learning models for MRI brain scan classification, developed using **Mojo 🔥** and **PyTorch**. |
|
|
|
|
|
**Full source code and training scripts:** [github.com/Meidverse/COMMRI](https://github.com/Meidverse/COMMRI) |
|
|
|
|
|
## Models |
|
|
|
|
|
### 1. Brain Tumor 2D CNN (`kaggle_tumor_2dcnn_best.pth`) |
|
|
|
|
|
| Metric | Value | |
|
|
|--------|-------| |
|
|
| **Accuracy** | 93.95% | |
|
|
| **Precision** | 0.94 | |
|
|
| **Recall** | 0.94 | |
|
|
| **F1 Score** | 0.94 | |
|
|
|
|
|
**Per-Class Performance:** |
|
|
| Class | Accuracy | |
|
|
|-------|----------| |
|
|
| Glioma | 98.1% | |
|
|
| Meningioma | 83.9% | |
|
|
| No Tumor | 98.5% | |
|
|
| Pituitary | 94.3% | |
|
|
|
|
|
### 2. IXI 3D Brain CNN (`ixi_3dcnn_best.pth`) |
|
|
|
|
|
3D CNN for brain volume classification from NIfTI files. |
|
|
|
|
|
## Quick Start |
|
|
|
|
|
### Option 1: Clone from GitHub (Recommended) |
|
|
|
|
|
```bash |
|
|
git clone https://github.com/Meidverse/COMMRI.git |
|
|
cd COMMRI |
|
|
|
|
|
# Install dependencies |
|
|
pip install -r requirements.txt |
|
|
|
|
|
# Run inference |
|
|
python -c " |
|
|
import torch |
|
|
from scripts.train_tumor import TumorCNN |
|
|
|
|
|
model = TumorCNN(4) |
|
|
model.load_state_dict(torch.load('kaggle_tumor_2dcnn_best.pth')) |
|
|
model.eval() |
|
|
print('Model loaded!') |
|
|
" |
|
|
``` |
|
|
|
|
|
### Option 2: Download from Hugging Face |
|
|
|
|
|
```python |
|
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
# Download model |
|
|
model_path = hf_hub_download( |
|
|
repo_id="Nikshey/mri-brain-classification", |
|
|
filename="kaggle_tumor_2dcnn_best.pth" |
|
|
) |
|
|
|
|
|
# Load with PyTorch |
|
|
import torch |
|
|
model = torch.load(model_path) |
|
|
``` |
|
|
|
|
|
## Inference Example |
|
|
|
|
|
```python |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
from torchvision import transforms |
|
|
from PIL import Image |
|
|
|
|
|
class TumorCNN(nn.Module): |
|
|
def __init__(self, num_classes=4): |
|
|
super().__init__() |
|
|
self.features = nn.Sequential( |
|
|
nn.Conv2d(3, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), |
|
|
nn.Conv2d(64, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), |
|
|
nn.MaxPool2d(2), nn.Dropout2d(0.25), |
|
|
nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(), |
|
|
nn.Conv2d(128, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(), |
|
|
nn.MaxPool2d(2), nn.Dropout2d(0.25), |
|
|
nn.Conv2d(128, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(), |
|
|
nn.Conv2d(256, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(), |
|
|
nn.MaxPool2d(2), nn.Dropout2d(0.25), |
|
|
nn.Conv2d(256, 512, 3, padding=1), nn.BatchNorm2d(512), nn.ReLU(), |
|
|
nn.Conv2d(512, 512, 3, padding=1), nn.BatchNorm2d(512), nn.ReLU(), |
|
|
nn.AdaptiveAvgPool2d(1), |
|
|
) |
|
|
self.classifier = nn.Sequential( |
|
|
nn.Flatten(), nn.Linear(512, 256), nn.ReLU(), nn.Dropout(0.5), |
|
|
nn.Linear(256, 128), nn.ReLU(), nn.Dropout(0.3), nn.Linear(128, num_classes), |
|
|
) |
|
|
def forward(self, x): |
|
|
return self.classifier(self.features(x)) |
|
|
|
|
|
# Load model |
|
|
model = TumorCNN(4) |
|
|
model.load_state_dict(torch.load("kaggle_tumor_2dcnn_best.pth", map_location="cpu")) |
|
|
model.eval() |
|
|
|
|
|
# Preprocess and predict |
|
|
transform = transforms.Compose([ |
|
|
transforms.Resize((224, 224)), |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) |
|
|
]) |
|
|
|
|
|
image = transform(Image.open("brain_mri.jpg").convert("RGB")).unsqueeze(0) |
|
|
pred = model(image).argmax(1).item() |
|
|
|
|
|
classes = ['glioma', 'meningioma', 'notumor', 'pituitary'] |
|
|
print(f"Prediction: {classes[pred]}") |
|
|
``` |
|
|
|
|
|
## Training |
|
|
|
|
|
Train your own models using the scripts in the [GitHub repo](https://github.com/Meidverse/COMMRI): |
|
|
|
|
|
```bash |
|
|
# Train tumor classifier |
|
|
mojo run scripts/train_tumor.mojo |
|
|
|
|
|
# Train 3D brain model |
|
|
mojo run scripts/train_advanced.mojo |
|
|
|
|
|
# Evaluate |
|
|
mojo run scripts/evaluate_tumor.mojo |
|
|
``` |
|
|
|
|
|
## Citation |
|
|
|
|
|
```bibtex |
|
|
@misc{commri2024, |
|
|
author = {Meidverse}, |
|
|
title = {COM-MRI: Brain Tumor Classification with Mojo}, |
|
|
year = {2024}, |
|
|
publisher = {GitHub}, |
|
|
url = {https://github.com/Meidverse/COMMRI} |
|
|
} |
|
|
``` |
|
|
|
|
|
## License |
|
|
|
|
|
MIT License - See [LICENSE](https://github.com/Meidverse/COMMRI/blob/main/LICENSE) |
|
|
|
|
|
## Acknowledgments |
|
|
|
|
|
- [Kaggle Brain Tumor MRI Dataset](https://www.kaggle.com/datasets/masoudnickparvar/brain-tumor-mri-dataset) |
|
|
- [IXI Dataset](https://brain-development.org/ixi-dataset/) |
|
|
- Built with [Mojo 🔥](https://www.modular.com/mojo) and PyTorch |
|
|
- Trained on NVIDIA RTX 4090 |
|
|
|