NeuroLens: Explainable MRI AI for Alzheimer’s Staging

Deep learning system that classifies Alzheimer’s stages from brain MRI scans using the Hack4Health dataset. We fine-tune a 2D/3D CNN with substantial data augmentation, apply Grad-CAM & error analysis

Comment

Inspiration

Alzheimer's disease affects millions globally, yet early and accurate diagnosis remains challenging. Medical imaging, particularly MRI scans, holds immense potential for detecting disease stages, but interpreting these scans requires significant expertise and time. We were inspired to develop NeuroLens to bridge this gap by providing an AI-powered tool that not only classifies Alzheimer's stages with high accuracy but also explains its decisions through Grad-CAM visualizations, making it trustworthy and useful for medical professionals.

What it does

NeuroLens is an explainable AI system that classifies Alzheimer's disease stages from brain MRI scans into 4 categories (stages 0-3). The system:

  • Processes and analyzes brain MRI images from the Hack4Health dataset (5,120 images)
  • Employs a custom 2D/3D Convolutional Neural Network (CNN) architecture with advanced regularization techniques
  • Achieves 61.72% test accuracy on multi-stage classification
  • Provides Grad-CAM (Gradient-weighted Class Activation Mapping) visualizations to highlight which regions of the brain influenced the model's predictions
  • Includes comprehensive error analysis and per-class performance metrics
  • Offers model interpretability crucial for medical AI applications

How we built it

Architecture & Technology Stack:

  • Deep Learning Framework: TensorFlow 2.19.0 with GPU acceleration
  • Model Architecture: Custom CNN with 3 convolutional blocks (32→64→128 filters)
  • Regularization: Batch Normalization, Dropout (0.25-0.5), and extensive Data Augmentation
  • Training Strategy: 70% train / 15% validation / 15% test split with stratified sampling
  • Optimizer: Adam with learning rate scheduling (0.001 → 6.25e-5)
  • Dataset: Alzheimer MRI Disease Classification Dataset from Kaggle (5,120 images)
  • Image Preprocessing: Grayscale conversion, 128x128 resize, normalization
  • Explainability: Grad-CAM implementation for model interpretability

Key Technical Decisions:

  • Data Augmentation: Rotation (15°), width/height shifts (0.1), zoom (0.1), horizontal flip
  • Callbacks: EarlyStopping (patience=10), ReduceLROnPlateau, ModelCheckpoint
  • Training: 41 epochs with automatic early stopping
  • Model saved as: neurolens_final_model.keras

Challenges we ran into

  1. Class Imbalance: The dataset had severe imbalance with only 49 samples for class 1 compared to 2,566 for class 2. This required careful stratified splitting and resulted in poor performance on minority classes.

  2. Model Overfitting: Initial models showed signs of overfitting with high training accuracy but poor validation performance. We addressed this through:

    • Multiple dropout layers (0.25-0.5)
    • Batch normalization after each convolutional and dense layer
    • Aggressive data augmentation
    • Learning rate reduction strategy

  3. Grad-CAM Implementation: Implementing Grad-CAM for Sequential models in TensorFlow 2.x required custom gradient tape operations and careful handling of layer outputs.

  4. Computational Resources: Training a deep CNN on 5,120 high-resolution MRI images required GPU acceleration and careful batch size management (batch_size=32).

Accomplishments that we're proud of

  • Achieved 61.72% test accuracy on a challenging 4-class medical imaging task with imbalanced data
  • Excellent performance on majority classes: 72.15% accuracy on stage 3 (267 samples) and 66.23% on stage 2 (385 samples)
  • Successfully implemented Grad-CAM visualizations for model interpretability, making the AI's decisions transparent and trustworthy for medical applications
  • Comprehensive evaluation framework with confusion matrices, per-class metrics, and training visualizations
  • Production-ready model saved and deployable (8.7M parameters, 33.23 MB)
  • Complete end-to-end pipeline from data loading to model deployment

What we learned

  • Medical AI requires explainability: Grad-CAM visualizations are crucial for building trust in medical AI systems
  • Class imbalance is a major challenge: Need better techniques like SMOTE, class weights, or focal loss for severe imbalance
  • Regularization is essential: Multiple regularization techniques (dropout, batch norm, augmentation) are needed for medical imaging
  • Transfer learning potential: Pre-trained models like ResNet or EfficientNet could potentially improve performance
  • Domain expertise matters: Medical imaging requires understanding of neuroanatomy and disease progression patterns

What's next for NeuroLens: Explainable MRI AI for Alzheimer's Staging

  1. Address class imbalance with advanced techniques (SMOTE, focal loss, class weights)
  2. Implement 3D CNN architecture to fully utilize volumetric MRI data
  3. Explore transfer learning with pre-trained medical imaging models
  4. Add attention mechanisms for better feature learning
  5. Expand to multi-modal inputs (MRI + clinical data + genetic markers)
  6. Clinical validation with real-world hospital data
  7. Deploy as web application for easy access by medical professionals
  8. Integrate with PACS systems for seamless clinical workflow integration

Built With

Share this project:

Updates