MRI Brain Tumor Classification Models
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
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)
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
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
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:
# 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
@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
Acknowledgments
- Kaggle Brain Tumor MRI Dataset
- IXI Dataset
- Built with Mojo ๐ฅ and PyTorch
- Trained on NVIDIA RTX 4090
- Downloads last month
- 10