Brain Tumor MRI Classification using Transfer Learning

A deep learning project for classifying brain tumor MRI images into 4 categories using transfer learning with EfficientNet-B0. This project was developed as part of an Applied Computer Vision course.

Table of Contents

• Project Overview
• Dataset
• Model Architecture
• Experiments
• Results
• Installation
• Usage
• Project Structure
• Team Members
• References

Project Overview

This project demonstrates transfer learning techniques for medical image classification. We fine-tune a pretrained EfficientNet-B0 model on brain tumor MRI scans to classify them into four categories:

1. Glioma - Tumors originating from glial cells in the brain
2. Meningioma - Tumors arising from the meninges (brain membranes)
3. Pituitary - Tumors in the pituitary gland
4. No Tumor - Healthy brain scans

Key Features

• Transfer learning with EfficientNet-B0 pretrained on ImageNet
• Zero-shot baseline evaluation to demonstrate transfer learning effectiveness
• Comprehensive experimentation with dropout regularization
• Data augmentation (horizontal flip, rotation)
• Mixed precision training for faster GPU computation
• Detailed visualizations and metrics analysis

Dataset

Brain Tumor MRI Dataset from Kaggle Source: masoudnickparvar/brain-tumor-mri-dataset

Split	Samples
Training	5,140
Validation	572 (10% of training)
Test	1,311
Total	7,023

Class Distribution

Class	Training	Test
Glioma	1,321	300
Meningioma	1,339	306
No Tumor	1,595	405
Pituitary	1,457	300

Model Architecture

We use EfficientNet-B0 as the backbone model, which provides an excellent balance between accuracy and computational efficiency.

• Pretrained weights: ImageNet-1K
• Input size: 224 x 224 pixels
• Parameters: ~4 million
• Libraries: timm, torchvision

Transfer Learning Strategy

1. Load pretrained EfficientNet-B0 weights
2. Replace classification head for 4-class output
3. Fine-tune entire network with Adam optimizer (lr=1e-4)
4. Apply regularization via dropout and stochastic depth

Experiments

We first established a zero-shot baseline (pretrained model with no training on tumor data) achieving ~42% accuracy, confirming the model has no prior knowledge of brain tumor classification. This demonstrates that our high accuracy results are genuinely learned through transfer learning.

We then conducted 4 experiments to analyze the effect of dropout on model generalization:

Experiment	Dropout	Drop Path	Epochs	Val Accuracy	Test Accuracy
1. Baseline (timm)	0.2	0.0	20	99.13%	99.08%
2. High Dropout	0.7	0.1	20	98.95%	98.55%
3. Custom (torchvision)	0.2	0.0	20	99.48%	99.47%
4. Extended Training	0.7	0.1	30	99.30%	98.86%

Best Model: Experiment 3 using torchvision's EfficientNet-B0 with default dropout (0.2) for 20 epochs.

Key Findings

1. Dropout variation had minimal impact - Our experiments revealed that adjusting dropout rates (0.2 vs 0.7) did not significantly affect final model performance. All configurations achieved >98.5% test accuracy, suggesting that the default dropout rate is already well-tuned for transfer learning scenarios.

2. Default dropout (0.2) is optimal - The standard dropout rate proved sufficient for this dataset. Higher dropout (0.7) actually showed slower convergence without improving generalization, indicating that aggressive regularization is unnecessary when fine-tuning pretrained models on this dataset size.

3. 20 epochs is sufficient - Extended training to 30 epochs showed no improvement over 20-epoch runs, suggesting early stopping is beneficial for efficiency.

4. Implementation consistency - Both timm and torchvision EfficientNet implementations achieved similar results, confirming the architecture's reliability across different libraries.

5. Transfer learning is highly effective - The jump from ~42% (zero-shot) to >99% (fine-tuned) demonstrates the power of transfer learning for medical image classification.

Results

Performance Metrics (Best Model)

• Test Accuracy: 99.47%
• Inference Time: ~1.1 ms/image (on RTX 3060 Laptop GPU)

Visualizations Generated

• Training/validation loss and accuracy curves
• Confusion matrix
• Per-class accuracy breakdown
• ROC curves (one-vs-rest)
• Precision-Recall curves
• Sample correct/incorrect predictions

Installation

Prerequisites

• Python 3.8+
• CUDA-capable GPU (recommended)
• Kaggle account (for dataset download)

Setup

1. Clone the repository

git clone https://github.com/ashkenazzio/brain-tumor-classification.git
cd brain-tumor-classification

2. Install dependencies

pip install -r requirements.txt

3. Get Kaggle API credentials

   • Go to Kaggle Settings
   • Click "Create New Token" to download kaggle.json
   • Either place kaggle.json in the project root directory, or enter credentials when prompted

Usage

Option 1: Google Colab (Recommended)

1. Open brain_tumor_classification.ipynb in Google Colab
2. Enable GPU runtime: Runtime > Change runtime type > GPU
3. Run the cells - you'll be prompted to enter your Kaggle credentials (or upload kaggle.json)
4. The notebook will download the dataset and run all experiments

Option 2: Local Jupyter

# Start Jupyter
jupyter notebook brain_tumor_classification.ipynb
# Enter Kaggle credentials when prompted (or place kaggle.json in project root)

Using the Trained Model

import torch
from torchvision import models, transforms
from PIL import Image

# Load model (Best model uses torchvision)
checkpoint = torch.load('checkpoints/efficientnet_b0_brain_tumor_best.pt')

model = models.efficientnet_b0(weights=None)
model.classifier = torch.nn.Sequential(
    torch.nn.Dropout(p=0.2),
    torch.nn.Linear(model.classifier[1].in_features, 4),
)
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

# Preprocess image
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Predict
image = Image.open('your_mri_image.jpg').convert('RGB')
input_tensor = transform(image).unsqueeze(0)

with torch.no_grad():
    output = model(input_tensor)
    prediction = output.argmax(1).item()

class_names = ['glioma', 'meningioma', 'notumor', 'pituitary']
print(f"Prediction: {class_names[prediction]}")

Project Structure

brain-tumor-classification/
├── brain_tumor_classification.ipynb  # Main notebook with all experiments
├── checkpoints/
│   └── efficientnet_b0_brain_tumor_best.pt  # Trained model
├── data/
│   └── .gitkeep                      # Dataset downloaded here at runtime
├── requirements.txt                  # Python dependencies
├── .gitignore
└── README.md

Team Members

• Deizy Lagziel (@Deizyy)
• Ofek Fuchs (@ofek-fuchs)
• Omri Ashkenazi (@ashkenazzio)
• Yuval Fischer (@Yuhan4)

References

1. EfficientNet: Tan, M., & Le, Q. (2019). EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks. ICML 2019.

2. Dataset: Nickparvar, M. (2021). Brain Tumor MRI Dataset. Kaggle. https://www.kaggle.com/datasets/masoudnickparvar/brain-tumor-mri-dataset

3. timm Library: Wightman, R. (2019). PyTorch Image Models. https://github.com/huggingface/pytorch-image-models

4. Transfer Learning: Yosinski, J., et al. (2014). How transferable are features in deep neural networks? NeurIPS 2014.

---

License

This project is licensed under the MIT License - see the LICENSE file for details.

---

This project was developed as part of the Applied Computer Vision course.