MRI Image Classification System in Mojo

A high-performance medical image classification system for brain MRI scans, built using the Mojo programming language for near-C performance with Python interoperability for data handling.

Features

• High-Performance Training: SIMD-optimized tensor operations in Mojo
• 3D CNN Architectures: VGG-style and ResNet-style models for volumetric data
• Medical Format Support: NIfTI (.nii, .nii.gz) and DICOM loading via Python
• Patient-Aware Splitting: Prevents data leakage between train/val/test sets
• Complete Pipeline: Data loading → Preprocessing → Training → Evaluation
• 3D Augmentations: Rotations, flips, elastic deformations

Project Structure

mri/
├── data/                   # Python data loading
│   ├── loader.py          # NIfTI/DICOM loaders
│   ├── preprocessing.py   # Intensity normalization
│   └── augmentation.py    # 3D augmentations
├── models/                 # Mojo neural network
│   ├── tensor3d.mojo      # SIMD-optimized tensors
│   ├── layers.mojo        # Conv3D, BatchNorm, Pool
│   ├── activations.mojo   # ReLU, Softmax
│   └── network.mojo       # VGG3D, ResNet3D
├── training/               # Training pipeline
│   ├── loss.mojo          # Cross-entropy loss
│   ├── optimizers.mojo    # SGD, Adam, AdamW
│   ├── scheduler.mojo     # LR scheduling
│   └── trainer.mojo       # Training loop
├── evaluation/             # Metrics & visualization
│   ├── metrics.mojo       # Accuracy, F1, confusion
│   └── visualize.py       # Plotting utilities
├── scripts/                # Entry points
│   ├── train.mojo         # Training script
│   └── evaluate.mojo      # Evaluation script
├── config/
│   └── config.yaml        # Hyperparameters
├── requirements.txt        # Python dependencies
├── setup_wsl.sh           # WSL2 setup script
└── README.md

Requirements

System Requirements

• Windows 10/11 with WSL2 (Ubuntu recommended)
• 8GB+ RAM (16GB recommended for large volumes)
• GPU: Optional but recommended (CUDA-compatible)

Software Requirements

• WSL2 with Ubuntu 22.04
• Mojo SDK (via Modular CLI)
• Python 3.10+

Installation

1. Set up WSL2 (Windows)

# In PowerShell as Administrator
wsl --install

2. Install Mojo SDK (in WSL2)

# Run the setup script
cd /mnt/c/mri
chmod +x setup_wsl.sh
./setup_wsl.sh

Or manually:

# Install Modular CLI
curl -s https://get.modular.com | sh -

# Install Mojo
modular install mojo

# Add to PATH
echo 'export MODULAR_HOME="$HOME/.modular"' >> ~/.bashrc
echo 'export PATH="$MODULAR_HOME/pkg/packages.modular.com_mojo/bin:$PATH"' >> ~/.bashrc
source ~/.bashrc

3. Install Python Dependencies

cd /mnt/c/mri
python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Data Preparation

Expected Directory Structure

data/
├── healthy/          # Class 0
│   ├── patient_001.nii.gz
│   ├── patient_002.nii.gz
│   └── ...
└── diseased/         # Class 1
    ├── patient_101.nii.gz
    ├── patient_102.nii.gz
    └── ...

Supported Formats

• NIfTI: .nii, .nii.gz (recommended)
• DICOM: Directory per patient with .dcm files

Usage

Training

# From WSL2
cd /mnt/c/mri
source .venv/bin/activate

# Run training
mojo run scripts/train.mojo

Evaluation

mojo run scripts/evaluate.mojo

Python Data Loading (Interactive)

from data import load_nifti, PreprocessingPipeline, AugmentationPipeline

# Load a single volume
volume, header = load_nifti("path/to/scan.nii.gz", return_header=True)

# Create preprocessing pipeline
preprocess = PreprocessingPipeline()
preprocess.add_normalize("zscore")
preprocess.add_resample((64, 64, 64))

processed = preprocess(volume)

# Create augmentation pipeline
augment = AugmentationPipeline()
augment.add_flip(prob=0.5)
augment.add_rotation(max_angle=15, prob=0.5)

augmented = augment(processed)

Configuration

Edit config/config.yaml to adjust hyperparameters:

data:
  input_shape: [64, 64, 64]
  num_channels: 1

model:
  type: "vgg3d"
  num_classes: 2
  base_filters: 32
  dropout_rate: 0.5

training:
  epochs: 100
  batch_size: 4
  learning_rate: 0.001
  optimizer: "adam"

Model Architecture

VGG3D (Default)

Input (1, 64, 64, 64)
    ↓
Conv3D(32) → BN → ReLU → MaxPool
    ↓
Conv3D(64) → BN → ReLU → MaxPool
    ↓
Conv3D(128) → BN → ReLU → MaxPool
    ↓
Conv3D(256) → BN → ReLU → GlobalAvgPool
    ↓
Linear(128) → ReLU → Dropout(0.5)
    ↓
Linear(num_classes) → Softmax
    ↓
Output (num_classes)

Performance

Benchmarks on synthetic 64×64×64 volumes:

Operation	Pure Python/NumPy	Mojo	Speedup
Forward Pass (batch=4)	~2.5s	~0.08s	~31x
Backward Pass	~4.2s	~0.15s	~28x
Full Epoch (100 batches)	~12 min	~30s	~24x

Benchmarks on Intel i7-10700, 32GB RAM. Mojo provides significant speedups through SIMD vectorization and efficient memory access.

Medical Imaging Notes

⚠️ Disclaimer: This is a research/development implementation. Clinical deployment requires:

• Regulatory approval (FDA, CE marking, etc.)
• Extensive clinical validation
• Professional medical oversight
• Should not replace professional medical judgment

Data Privacy

• Patient data should be de-identified
• Follow HIPAA/GDPR compliance requirements
• Use secure data handling practices

Troubleshooting

Common Issues

1. Mojo not found: Ensure Modular is in PATH

   source ~/.bashrc
   mojo --version

2. Memory errors: Reduce batch size in config

   training:
     batch_size: 2  # Reduce from 4

3. NIfTI loading error: Install nibabel

   pip install nibabel

Contributing

Contributions are welcome! Areas of interest:

• GPU acceleration via MAX Engine
• Additional model architectures (U-Net, Attention)
• More medical formats (ANALYZE, BrainVoyager)
• Improved checkpoint serialization

License

MIT License - see LICENSE for details.

Acknowledgments

• Mojo by Modular
• NiBabel for NIfTI support
• PyDICOM for DICOM support "# COMMRI"