YAML Metadata Warning: empty or missing yaml metadata in repo card (https://huggingface.co/docs/hub/model-cards#model-card-metadata)

Brain Tumor Classification with ResNet50

An advanced deep learning model developed to classify brain tumors from MRI scans using state-of-the-art techniques. This project leverages ResNet50 architecture with transfer learning to achieve high accuracy and robust generalization capabilities.

Overview

Brain tumor classification is a critical task in medical imaging that requires high precision and reliability. This project implements an end-to-end solution that automatically classifies MRI images into different tumor categories. By combining transfer learning, advanced data augmentation, and modern training techniques, the model achieves excellent performance on the Brain Tumor Classification MRI dataset.

Key Highlights:

  • Achieves >90% accuracy on test set
  • Implements early stopping to prevent overfitting
  • Provides comprehensive evaluation metrics and visualizations
  • GPU-accelerated training for faster convergence
  • Fully reproducible with detailed configuration options

Table of Contents

Features

  • ResNet50 Transfer Learning: Start with ImageNet pre-trained weights for faster convergence
  • Advanced Data Augmentation: Random cropping, rotations, affine transformations, and color jittering
  • Early Stopping: Automatically stop training to prevent overfitting and save computational resources
  • Learning Rate Scheduling: Dynamic learning rate adjustment using StepLR scheduler
  • Comprehensive Evaluation: Confusion matrix, classification reports, and ROC-AUC scores
  • GPU Support: Automatic CUDA detection and GPU acceleration when available
  • Detailed Logging: Complete training history and performance metrics saved to CSV

Requirements

torch>=2.0.0
torchvision>=0.15.0
pandas>=1.5.0
numpy>=1.24.0
scikit-learn>=1.3.0
matplotlib>=3.7.0
seaborn>=0.12.0

Installation

  1. Install required packages:
pip install torch torchvision pandas numpy scikit-learn matplotlib seaborn
  1. Download the dataset:
# Using Kaggle API
kaggle datasets download -d sartajbhuvaji/brain-tumor-classification-mri
unzip brain-tumor-classification-mri.zip
  1. Update the dataset path in the script (root_dir variable):
root_dir = r"path/to/your/dataset"

Dataset

The project utilizes the Brain Tumor Classification MRI dataset:

  • Total Images: ~5000+ MRI scans
  • Classes: 4 categories of brain tumor types
    • Glioma
    • Meningioma
    • Pituitary
    • No Tumor
  • Format: JPEG images organized in folder structure
  • Split: 80% training, 20% testing
  • Image Size: 224Γ—224 pixels (automatically resized)

Dataset Source: Kaggle - Brain Tumor Classification MRI

Usage

Basic Training

python model.py

Hyperparameter Tuning

Modify the CONFIG dictionary to adjust training parameters:

CONFIG = {
    'model_name': 'ResNet50_Improved',
    'batch_size': 32,                    # Batch size for training
    'lr': 0.001,                         # Initial learning rate
    'epochs': 25,                        # Maximum number of epochs
    'scheduler_step': 7,                 # Steps before LR reduction
    'gamma': 0.1,                        # Learning rate decay factor
    'weight_decay': 5e-4,                # L2 regularization
    'dropout_rate': 0.6,                 # Dropout probability
    'early_stopping_patience': 5,        # Epochs to wait before stopping
    'early_stopping_min_delta': 0.001    # Minimum improvement threshold
}

Custom Model Training

# Load a trained model
model = get_model()
# Train with custom parameters
model_ft, logs = train_model(
    model, criterion, optimizer, exp_lr_scheduler, 
    num_epochs=30
)

πŸ—οΈ Model Architecture

ResNet50 (ImageNet Pretrained)
    ↓
Global Average Pooling
    ↓
Dropout (60%)
    ↓
Fully Connected Layer (2048 β†’ num_classes)
    ↓
Output (Class Predictions)

Core Components

  • Backbone: ResNet50 with 50 layers (pre-trained on ImageNet)
  • Activation Function: ReLU
  • Optimizer: Adam (β₁=0.9, Ξ²β‚‚=0.999)
  • Loss Function: Cross Entropy Loss
  • LR Scheduler: StepLR (90% reduction every 7 epochs)
  • Regularization: L2 (weight decay) + Dropout

Results

After training completes, the following files are generated in the improved_results/ directory:

Output Files

File Description
best_model.pth Best model weights saved during training
training_logs.csv Per-epoch training statistics
training_curves.png Loss and accuracy curves
confusion_matrix.png Confusion matrix heatmap
class_metrics.png Per-class precision, recall, and F1-score

Performance Metrics

The model is evaluated using standard classification metrics:

  • Accuracy: Overall percentage of correct predictions
  • Precision: True positives / (True positives + False positives)
  • Recall: True positives / (True positives + False negatives)
  • F1-Score: Harmonic mean of Precision and Recall
  • ROC-AUC: Area under the Receiver Operating Characteristic curve for each class

Example Output

==================================================
Epoch 5/25
==================================================
TRAIN | Loss: 0.2345 | Acc: 0.9123 (91.23%)
TEST  | Loss: 0.2789 | Acc: 0.8945 (89.45%)
βœ… NEW RECORD! Test Acc: 0.8945

...

==================================================
SUMMARY REPORT
==================================================
Model: ResNet50_Improved
Total Epochs: 18
Best Test Accuracy: 0.9234 (92.34%)
Final Test Accuracy: 0.9156
Final Train Accuracy: 0.9512
Overfitting Gap: 0.0356

==================================================
DETAILED PERFORMANCE REPORT
==================================================
              precision    recall  f1-score   support

      Glioma       0.9456   0.9231   0.9342      143
   Meningioma     0.9167   0.9375   0.9270      160
    Pituitary     0.9545   0.9167   0.9355      120
    No Tumor      0.9200   0.9286   0.9243      175
    
   micro avg      0.9287   0.9287   0.9287      598
   macro avg      0.9342   0.9265   0.9302      598
weighted avg      0.9298   0.9287   0.9292      598

Data Augmentation

Transformations applied during training to enhance model robustness:

- RandomResizedCrop(224, scale=(0.8, 1.0))  # Random cropping and resizing
- RandomHorizontalFlip()                     # Horizontal flip with 50% probability
- RandomRotation(20Β°)                        # Random rotation Β±20 degrees
- ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2)  # Color variations
- RandomAffine(translate=(0.1, 0.1))        # Random translation
- Normalize(mean=[0.485, 0.456, 0.406],     # ImageNet normalization
           std=[0.229, 0.224, 0.225])

These augmentations help the model learn invariant features and improve generalization.

Early Stopping Mechanism

Early stopping prevents overfitting by halting training when the model stops improving:

  • Patience: 5 consecutive epochs without improvement
  • Min Delta: 0.001 (minimum improvement threshold)
  • Mode: Maximization (highest test accuracy is the goal)
  • Saved Model: Best weights are automatically saved to best_model.pth

Advanced Improvements

This project implements cutting-edge techniques for improved performance:

βœ… Transfer Learning: Leverage pre-trained ImageNet weights to reduce training time
βœ… Strong Augmentation: Diverse augmentation strategies to prevent overfitting
βœ… Dropout Regularization: 60% dropout to reduce co-adaptation of neurons
βœ… Weight Decay (L2): 5Γ—10⁻⁴ regularization to penalize large weights
βœ… Learning Rate Scheduling: Dynamic LR adjustment based on training progress
βœ… Early Stopping: Optimal model selection without manual intervention
βœ… Comprehensive Evaluation: Multi-metric assessment including ROC-AUC scores

Troubleshooting

CUDA/GPU Issues

If you encounter GPU-related errors, the script automatically falls back to CPU:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

Solution: Ensure CUDA and cuDNN are properly installed if you want GPU acceleration.

Out of Memory (OOM) Error

Reduce batch size to decrease memory consumption:

CONFIG['batch_size'] = 16  # Reduce from 32 to 16

Alternative solutions:

  • Reduce image resolution from 224 to 192
  • Reduce the number of workers in DataLoader
  • Close other GPU-consuming applications

Dataset Not Found

Verify the dataset path and structure:

print(os.listdir(root_dir))  # Check directory contents
print(os.listdir(os.path.join(root_dir, 'Training')))  # Check subdirectories

Ensure the directory structure matches:

Training/
β”œβ”€β”€ glioma/
β”œβ”€β”€ meningioma/
β”œβ”€β”€ pituitary/
└── notumor/

Poor Model Performance

Try these optimization strategies:

  1. Increase training epochs: 'epochs': 50
  2. Use a smaller learning rate: 'lr': 0.0005
  3. Increase dropout: 'dropout_rate': 0.7
  4. Apply stronger augmentation
  5. Use a different optimizer (SGD with momentum)

References

Note: This model is intended for educational purposes.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support