COM-MRI

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)

Per-Class Performance:

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