Brain Tumor MRI Classification
Train a ResNet-50 model to classify brain MRI scans and deploy it as a web app.
Introduction
30 Second Summary
When a radiologist examines a brain MRI, they are looking for subtle differences in shape, texture, and location to determine whether a mass is dangerous. That process takes years of training and careful attention to detail.
In this project, you will train a deep learning model to classify brain MRI scans into four tumor categories, then deploy it as a live web app where anyone can upload an image and get an instant prediction.
What You'll Build
You will upload a brain MRI image to your app, and within a second it returns a prediction with confidence percentages for each of the four tumor types.
By the end of this project, you'll have:
- A trained ResNet-50 model that classifies brain MRIs into Glioma, Meningioma, Pituitary, or No Tumor with over 90% accuracy.
- A live Gradio web app with a shareable public URL where anyone can upload an MRI scan and see confidence scores across all four classes.
- A GitHub repository containing your full training notebook and deployment code, ready for permanent hosting on HuggingFace Spaces.
- Secret Mission: Add Grad-CAM visualization so your app shows a heatmap of where in the MRI the model is looking to make its decision.
Are there any prerequisites?
You should be comfortable with basic Python and have a general familiarity with neural networks (what layers are, how training works). No prior PyTorch or image classification experience is needed.
Before We Start
Before we dive in, let's lock in what you're about to build and why it matters.
Taking a moment to connect the project goal to your own learning journey helps you stay focused through the technical steps ahead.
Set Up Your Environment and Download the Dataset
Training a deep learning model on medical images requires two things: a GPU to handle the heavy computation and a labeled dataset to learn from.
In this step, you will set up a free GPU-powered notebook in Google Colab and download a brain tumor MRI dataset from Kaggle. By the end, you will have everything loaded and ready for model training.
In this step, get ready to:
- Open Google Colab and set the runtime to a T4 GPU.
- Download the brain tumor MRI dataset using the Kaggle CLI.
- Verify the dataset structure has 4 tumor classes across train, validation, and test splits.
Open Google Colab and configure the GPU runtime
Google Colab gives you a free notebook environment with GPU access. The T4 GPU will speed up model training dramatically compared to CPU-only execution.
- Open your browser and navigate to colab.research.google.com.
- Sign in with your Google account if prompted.
- Click New notebook to create a fresh notebook.
- In the menu bar, click Runtime.
- Click Change runtime type.
- In the dialog that appears, set Hardware accelerator to T4 GPU.
- Click Save.
What is Google Colab?
Google Colab is a free cloud-based Jupyter notebook environment. It comes with Python 3.12 and PyTorch 2.8.0 pre-installed, so you do not need to set up anything locally.
The T4 GPU Colab provides for free is the same hardware used in production machine learning workloads. It will cut training time from hours (on CPU) down to minutes.
You should now see an empty notebook with a single code cell, ready for input.
Don't see the T4 GPU option?
Colab's free GPU availability varies by demand. If T4 is not listed, try again in a few minutes or check that you are not already running another GPU session.
Help me find the T4 GPU option in Colab.
Upload your Kaggle credentials and download the dataset
The brain tumor MRI dataset lives on Kaggle, a platform for datasets and machine learning competitions. To download it programmatically, you need an API token from your Kaggle account.
- In your browser, open a new tab and go to kaggle.com/settings/api.
- Click Create New Token.
- A file called kaggle.json downloads to your computer. This contains your API credentials.
What is kaggle.json?
This file contains your Kaggle username and a secret API key. The Kaggle CLI reads it to authenticate your download requests. Never share this file or commit it to a public repository.
- Switch back to your Colab notebook.
- Upload your kaggle.json file to the Colab session by running this code in the first cell:
# Upload your kaggle.json credentials file to this Colab session
from google.colab import files
files.upload()
What does this code do?
The files.upload() function opens a file picker dialog in your browser. Select the kaggle.json file you just downloaded from Kaggle.
After clicking the cell's play button, a file picker appears below the cell. Select your kaggle.json file and wait for the upload confirmation message.
- Set up the credentials and download the dataset by running this code in a new cell:
# Install the Kaggle CLI
!pip install kaggle
# Move credentials to the location Kaggle expects
!mkdir -p ~/.kaggle
!mv kaggle.json ~/.kaggle/
!chmod 600 ~/.kaggle/kaggle.json
# Download and extract the brain tumor MRI dataset
!kaggle datasets download -d mohiburrahmanrifat/brain-tumor-mri --unzip -p brain-tumor-mri
What does this code do?
- pip install kaggle installs the Kaggle command-line tool so you can download datasets programmatically.
- The mkdir, mv, and chmod commands place your credentials where the Kaggle CLI expects to find them and restrict file permissions for security.
- The final command downloads and extracts the dataset into a folder called brain-tumor-mri/. The --unzip flag automatically extracts the archive after download.
You will see installation output followed by download progress. The download is approximately 150 MB and takes about 30 seconds on Colab's connection.
Seeing a 403 Forbidden or credentials error?
Make sure you uploaded the correct kaggle.json file in the previous cell. If you accidentally uploaded a different file, re-run the files.upload() cell and select the right one.
Also confirm your Kaggle account has accepted the dataset's terms of use by visiting the dataset page directly.
Help me fix my Kaggle API credentials error in Google Colab.
Verify your GPU and dataset structure
Before moving on, you need to confirm two things: the GPU is active and the dataset extracted correctly.
Before you run the next command, what do you expect to see if the GPU is set up correctly?
- Check that your GPU is active by running this in a new cell:
# Confirm the T4 GPU is available for training
!nvidia-smi
You should see a table showing a Tesla T4 GPU with available memory. This confirms Colab assigned GPU hardware to your session.
- Now check the dataset folders by running this in a new cell:
# List the class folders inside the Training split
!ls brain-tumor-mri/Training/
You should see four subdirectories listed: Glioma, Meningioma, No Tumor, and Pituitary.
What are these four classes?
Each subdirectory represents one type of brain MRI classification your model will learn. Glioma, Meningioma, and Pituitary are three common brain tumor types. The No Tumor class contains healthy brain scans for comparison.
The dataset contains 7,200 images total, pre-split into Training (5,600), Validation (800), and Testing (800) sets. Each image is a 224x224 pixel JPG file.
- Confirm the Validation and Testing splits have the same structure by running:
# Verify all three splits have the same 4-class structure
!ls brain-tumor-mri/Validation/
!ls brain-tumor-mri/Testing/
Both commands should output the same four class directories: Glioma, Meningioma, No Tumor, Pituitary.
Seeing 'No such file or directory'?
This means the download did not complete or extracted to a different location. Run !ls in a cell to see what files exist in your current directory.
If you see a .zip file instead of a folder, the --unzip flag may not have worked. Re-run the download command or manually unzip with !unzip.
Help me fix my dataset download issue in Colab.
Your environment is ready with a T4 GPU and 7,200 labeled MRI images across 4 classes. Next up, you will load this data into PyTorch and train your first baseline model.
Load the Data and Train a Baseline Model
Your environment is ready and the dataset is sitting in your Colab storage. Now you need to answer a fundamental question: how well can a powerful pretrained model perform on brain MRI classification with minimal effort?
In this step, you will load the MRI images into PyTorch, set up a pretrained ResNet-50 model for transfer learning, and train it to classify brain tumors. This establishes a baseline accuracy number that you will improve in later steps.
In this step, get ready to:
- Load the MRI dataset using PyTorch ImageFolder with ImageNet normalization.
- Configure ResNet-50 for 4-class classification with frozen backbone layers.
- Train for 5 epochs and evaluate baseline test accuracy.
Load the MRI images with PyTorch ImageFolder
PyTorch's ImageFolder reads images from a directory where each subdirectory name becomes a class label. Your dataset already has this exact structure.
- Create a new code cell in your Colab notebook and paste the following code to load the dataset with basic transforms:
import torch
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
# Define transforms: resize to 224x224, convert to tensor, normalize with ImageNet stats
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])
])
# Load each split using ImageFolder (directory names become class labels)
train_dataset = ImageFolder('brain-tumor-mri/Training', transform=transform)
val_dataset = ImageFolder('brain-tumor-mri/Validation', transform=transform)
test_dataset = ImageFolder('brain-tumor-mri/Testing', transform=transform)
# Create DataLoaders for batching and shuffling
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
# Verify the dataset loaded correctly
print(f"Classes: {train_dataset.classes}")
print(f"Training samples: {len(train_dataset)}")
print(f"Validation samples: {len(val_dataset)}")
print(f"Test samples: {len(test_dataset)}")
What does this code do?
- transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) applies the same normalization that ResNet-50 was trained with on ImageNet. This ensures the pixel values match what the model expects.
- ImageFolder automatically assigns numeric labels based on alphabetically sorted subdirectory names (Glioma=0, Meningioma=1, No Tumor=2, Pituitary=3).
- DataLoader wraps the dataset to serve images in batches of 32, which is more efficient than loading one image at a time.
- Run the cell. You should see the 4 class names printed along with sample counts (Training: 5,600, Validation: 800, Testing: 800).
Seeing a FileNotFoundError?
Make sure your dataset path matches exactly. Run !ls brain-tumor-mri/ to confirm the folder exists. If you see Training, Validation, and Testing listed, the path is correct.
help me fix a FileNotFoundError when loading my brain tumor MRI dataset with PyTorch ImageFolder in Google Colab
Load ResNet-50 and configure it for transfer learning
ResNet-50 is a CNN pretrained on 1.2 million images from ImageNet. Instead of training a model from scratch with only 5,600 MRI images, you will reuse ResNet-50's learned features and only replace its final classification layer.
- Create a new code cell and paste this code to load and configure the model:
from torchvision.models import resnet50, ResNet50_Weights
import torch.nn as nn
# Load ResNet-50 pretrained on ImageNet (80.9% top-1 accuracy)
model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
# Freeze all layers so their weights don't change during training
for param in model.parameters():
param.requires_grad = False
# Replace the final fully connected layer (originally 1000 classes) with 4 classes
model.fc = nn.Linear(model.fc.in_features, 4)
# Move model to GPU for faster training
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
print(f"Using device: {device}")
print(f"Output layer: {model.fc}")
Why freeze the backbone layers?
Freezing means the pretrained weights stay fixed during training. Only the new final layer (4 outputs) learns from your MRI data. This is much faster than updating all 25 million parameters and works well when your dataset is small.
model.fc.in_features grabs the size of the layer feeding into the classifier (2048 for ResNet-50). Your new layer maps those 2048 features to 4 tumor classes.
- Run the cell. You should see Using device: cuda and Output layer: Linear(in_features=2048, out_features=4, bias=True) printed.
Seeing 'cpu' instead of 'cuda'?
Your runtime is not set to GPU. Go to **Runtime** in the top menu, click **Change runtime type**, and select **T4 GPU**. Then restart the runtime and re-run all cells from the top.
help me switch my Google Colab runtime to use a T4 GPU
Train for 5 epochs and evaluate baseline accuracy
Now you will train the model. Since only the final layer is unfrozen, training 5 epochs should take about 2-3 minutes on the T4 GPU.
- Create a new code cell and paste the training loop:
import torch.optim as optim
# Only optimize the new classification head (frozen layers are excluded)
optimizer = optim.Adam(model.fc.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
train_losses, val_losses = [], []
for epoch in range(5):
model.train()
running_loss, correct, total = 0.0, 0, 0
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
train_acc = 100 * correct / total
train_losses.append(running_loss / len(train_loader))
# Validation phase (no gradient computation needed)
model.eval()
val_loss, val_correct, val_total = 0.0, 0, 0
with torch.no_grad():
for images, labels in val_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
val_loss += loss.item()
_, predicted = torch.max(outputs, 1)
val_total += labels.size(0)
val_correct += (predicted == labels).sum().item()
val_acc = 100 * val_correct / val_total
val_losses.append(val_loss / len(val_loader))
print(f"Epoch {epoch+1}/5 - Train Acc: {train_acc:.1f}%, Val Acc: {val_acc:.1f}%")
What does this training loop do?
- Each epoch loops through all training batches, computes predictions, measures the error (cross-entropy loss), and updates the final layer's weights using the Adam optimizer.
- The validation phase runs after each epoch with torch.no_grad() to measure performance on unseen data without updating weights.
- You will see accuracy increase each epoch as the final layer learns to map ResNet-50's features to tumor classes.
- Run the cell. You will see accuracy printed for each epoch, with training taking about 2-3 minutes total.
Training seems stuck or very slow?
Confirm you are on a GPU runtime by running !nvidia-smi in a new cell. If it shows a Tesla T4, training is running. Each epoch processes 175 batches, so give it about 30 seconds per epoch.
help me debug slow PyTorch training in Google Colab
- Now plot the loss curves to visualize training progress. Create a new code cell:
import matplotlib.pyplot as plt
# Plot training vs validation loss over epochs
plt.figure(figsize=(8, 4))
plt.plot(range(1, 6), train_losses, label='Train Loss')
plt.plot(range(1, 6), val_losses, label='Val Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training vs Validation Loss')
plt.legend()
plt.show()
- Run the cell. You should see a line chart with both train and validation loss decreasing over the 5 epochs.
Before you evaluate on the test set, what accuracy do you expect from a model that only uses its final layer to classify MRI scans? Think about whether 5 epochs with a frozen backbone is enough for medical images.
- Create a new code cell to evaluate on the test set:
# Evaluate on the held-out test set
model.eval()
test_correct = 0
test_total = 0
with torch.no_grad():
for images, labels in test_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs, 1)
test_total += labels.size(0)
test_correct += (predicted == labels).sum().item()
test_accuracy = 100 * test_correct / test_total
print(f"Baseline Test Accuracy: {test_accuracy:.2f}%")
What should I see?
Your baseline test accuracy should land between 85-92%. This is impressive for just 5 epochs of training a single layer. But a single accuracy number hides important details about which tumor types the model confuses.
- Run the cell. Note your test accuracy number. You will use this as the benchmark to beat after adding improvements.
✔️ Awesome, I've got everything!
Great. Make sure your notebook has all 4 cells from this step saved and each produced output.
ⓧ I'd like to double check the full code
Scroll through your notebook and confirm these cells exist in order:
- Cell 1: Imports, transforms, ImageFolder datasets, DataLoaders, and print statements showing classes and sample counts.
- Cell 2: ResNet-50 model loading, layer freezing, FC replacement, device assignment, and print showing cuda and output layer info.
- Cell 3: Training loop for 5 epochs printing train/val accuracy each epoch.
- Cell 4: Loss curve plot (matplotlib chart visible).
- Cell 5: Test evaluation printing baseline accuracy between 85-92%.
You have a working baseline model hitting 85-92% accuracy with minimal effort. That is the power of transfer learning. But which tumor types is it getting wrong? In the next step, you will dig into the confusion matrix and discover exactly where the model struggles.
Evaluate and Diagnose the Model
Your baseline model is trained and you have a test accuracy number in hand. That single number tells you the model works, but it hides a critical question: which classes is the model getting wrong?
In this step, you'll look beyond overall accuracy and examine per-class performance. A confusion matrix will reveal exactly where your model struggles, exposing which tumor types it confuses with each other.
In this step, get ready to:
- Run inference on the full test set and collect all predictions.
- Generate and plot a confusion matrix heatmap showing misclassifications per class.
- Print a classification report and identify the weakest class pair.
Run inference on the full test set
To understand per-class performance, you need the model's prediction for every single test image alongside the true label. Collecting all of these into lists lets you feed them into evaluation tools that break down accuracy by class.
- In a new Colab cell, run the model on every test batch and collect predictions by running this code:
import numpy as np
# Collect all predictions and ground truth labels
all_preds = []
all_labels = []
model.eval()
with torch.no_grad():
for images, labels in test_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs, 1)
all_preds.extend(predicted.cpu().numpy())
all_labels.extend(labels.cpu().numpy())
print(f"Total test samples evaluated: {len(all_preds)}")
What does this code do?
model.eval() switches the model out of training mode so batch normalization and dropout behave correctly during inference.
torch.no_grad() disables gradient computation, which speeds up inference and saves memory since you are not updating weights. Each batch of images goes through the model, and torch.max picks the highest-confidence class as the prediction.
You should see output like Total test samples evaluated: 800 confirming all test images were processed.
Seeing 0 samples or an error?
Make sure test_loader was defined in a previous cell and that cell was executed. If you restarted the Colab runtime, you need to re-run all previous cells from the top.
help me debug why my test_loader shows 0 samples in Colab
Generate and plot the confusion matrix
A confusion matrix is a grid where rows represent the actual class and columns represent what the model predicted. Off-diagonal cells show misclassifications, so a bright off-diagonal spot means the model consistently confuses those two classes.
Before you run the next cell, think about which two tumor classes you expect the model to confuse the most. Consider that Glioma and Meningioma can look visually similar in MRI scans.
- Plot the confusion matrix as an annotated heatmap by running this code in a new cell:
from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
# Get class names from the dataset
class_names = test_dataset.classes
# Compute confusion matrix from predictions vs ground truth
cm = confusion_matrix(all_labels, all_preds)
# Plot as an annotated heatmap
plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
xticklabels=class_names, yticklabels=class_names)
plt.xlabel('Predicted')
plt.ylabel('Actual')
plt.title('Confusion Matrix - Baseline Model')
plt.tight_layout()
plt.show()
What does this code do?
confusion_matrix(all_labels, all_preds) computes a 4x4 grid where each cell counts how many times the model predicted class j when the true class was i.
sns.heatmap with annot=True and fmt='d' overlays the raw integer counts on each cell so you can read exact numbers. Darker blue cells mean higher counts.
You should see a 4x4 heatmap appear below the cell. The diagonal cells (top-left to bottom-right) show correct predictions.
Look at the off-diagonal cells. The largest off-diagonal numbers appear between Glioma and Meningioma, confirming these two classes are the ones your model confuses most. This is expected. The frozen backbone lacks the sensitivity to distinguish between tumors with similar visual appearance in MRI.
Not seeing the heatmap?
If the plot does not render, check that you have import matplotlib.pyplot as plt and import seaborn as sns at the top of the cell. Both are pre-installed in Colab but must be explicitly imported.
If you see a NameError for test_dataset, make sure the cell where you defined your test dataset with ImageFolder was run.
help me fix my confusion matrix plot in Colab
Print the classification report
The confusion matrix gives a visual overview, but you also want exact numbers. A classification report breaks down precision, recall, and F1-score for each class individually.
- Print the full per-class breakdown by running this code in a new cell:
from sklearn.metrics import classification_report
# Print precision, recall, and F1-score for each class
report = classification_report(all_labels, all_preds,
target_names=class_names)
print(report)
What do these metrics mean?
Precision tells you: of all images the model labeled as class X, what fraction actually were class X.
Recall tells you: of all images that truly are class X, what fraction did the model catch.
F1-score is the harmonic mean of precision and recall. A low F1 for a class means the model is weak there, either missing real cases or generating false positives.
Look at the report output. Glioma and Meningioma have the lowest F1-scores compared to No Tumor and Pituitary. This confirms what the confusion matrix showed visually.
Why does the model confuse Glioma and Meningioma?
Both tumor types appear as masses within similar brain regions on T1-weighted MRI scans. With the backbone frozen, the model only uses generic ImageNet features (edges, textures, shapes) rather than learning fine-grained medical patterns.
The next step addresses this by unfreezing deeper layers and adding augmentation so the model can learn the subtle differences between these two tumor types.
✔️ Awesome, I've got everything!
Great. Make sure you can clearly see the Glioma-Meningioma confusion in your heatmap before moving on.
ⓧ I'd like to double check the full code
Scroll up through this step and confirm you have these three code chunks in order:
- Test inference loop (collects all_preds and all_labels from test_loader).
- Confusion matrix heatmap (uses sklearn.metrics.confusion_matrix and seaborn to plot a 4x4 grid).
- Classification report (uses sklearn.metrics.classification_report with target_names=class_names).
Each chunk should be in its own Colab cell. The output should show 800 test samples, a visible heatmap, and a printed report.
You have now diagnosed the baseline model's blind spot. Next up, you'll fix it with data augmentation and deeper fine-tuning to push past this Glioma-Meningioma confusion.
Improve with Data Augmentation and Fine-Tuning
Your confusion matrix told a clear story: the frozen backbone can't tell Glioma apart from Meningioma. These two tumor types look similar in MRI scans, and generic ImageNet features aren't sensitive enough to pick up the subtle differences.
In this step, you'll fix that shortfall with two techniques. First, data augmentation exposes the model to rotated, flipped, and brightness-shifted versions of each image so it can't memorize surface-level patterns. Second, unfreezing deeper layers lets the model learn MRI-specific features instead of relying solely on ImageNet representations.
In this step, get ready to:
- Add medical-image-specific data augmentation to training transforms.
- Unfreeze deeper ResNet-50 layers with differential learning rates.
- Retrain for 10 epochs and save the best model checkpoint.
Add data augmentation to the training transforms
Data augmentation artificially expands your training set by applying random transformations to each image every time it's loaded. The model sees a slightly different version each epoch, which prevents overfitting and forces it to learn robust features.
- Create a new augmented transform pipeline and rebuild your training dataloader by running this code in a new cell:
# Augmented transforms: force the model to learn robust features
train_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(15),
transforms.ColorJitter(brightness=0.2, contrast=0.2),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# Recreate training dataset and loader with augmented transforms
train_dataset = datasets.ImageFolder('brain-tumor-mri/Training', transform=train_transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
# Confirm the augmented data loads correctly
images, labels = next(iter(train_loader))
print(f'Batch shape: {images.shape}')
print(f'Training samples: {len(train_dataset)}')
What does this code do?
- RandomHorizontalFlip() mirrors each image with 50% probability, simulating the fact that tumors can appear on either side of the brain.
- RandomRotation(15) rotates images up to 15 degrees, handling slight variations in how MRI scans are positioned.
- ColorJitter(brightness=0.2, contrast=0.2) randomly shifts brightness and contrast, making the model robust to differences in MRI machine calibration.
You should see output confirming a batch shape of torch.Size([32, 3, 224, 224]) and 5,600 training samples.
Seeing a file not found error?
- Make sure the path brain-tumor-mri/Training matches the directory where you unzipped the dataset in Step 1.
- If your Colab runtime disconnected, the dataset files were deleted. Re-run your download cell from Step 1 first.
Help me fix a file path error in my Colab notebook when loading the brain tumor MRI dataset.
Unfreeze layers and configure differential learning rates
With the backbone frozen, only the final classifier layer learns. To capture subtle MRI-specific details, you need to let deeper convolutional layers adapt too. You'll use a lower learning rate for these pretrained layers so they refine rather than destroy what they already know.
- Unfreeze the deeper ResNet-50 layers and set up the optimizer with differential learning rates by running this cell:
# Unfreeze layer3 and layer4 for fine-tuning
for param in model.layer3.parameters():
param.requires_grad = True
for param in model.layer4.parameters():
param.requires_grad = True
# Differential learning rates: lower for pretrained layers, higher for classifier
optimizer = torch.optim.Adam([
{'params': model.layer3.parameters(), 'lr': 1e-4},
{'params': model.layer4.parameters(), 'lr': 1e-4},
{'params': model.fc.parameters(), 'lr': 1e-3}
])
# Learning rate scheduler: reduce by 10x every 3 epochs
from torch.optim.lr_scheduler import StepLR
scheduler = StepLR(optimizer, step_size=3, gamma=0.1)
# Check how many parameters are now trainable
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f'Trainable parameters: {trainable:,} / {total:,}')
What does this code do?
- Setting requires_grad = True on layer3 and layer4 allows these deeper residual blocks to update during backpropagation.
- The optimizer uses parameter groups: backbone layers train at 1e-4 (gentle refinement) while the classifier head trains at 1e-3 (faster adaptation to the new 4-class task).
- StepLR multiplies all learning rates by 0.1 every 3 epochs, letting training converge smoothly as accuracy plateaus.
You should see the trainable parameter count jump significantly compared to the baseline (where only the classifier head was trainable).
Still seeing the same parameter count?
- Make sure you ran the cell. In Colab, pressing Shift+Enter executes the current cell.
- Confirm model.layer3 and model.layer4 exist. If you renamed your model variable, update the code accordingly.
Help me debug why my trainable parameter count didn't change after unfreezing layers.
Retrain and evaluate the improved model
With augmented data flowing in and deeper layers unfrozen, you're ready to retrain. The training loop now saves a checkpoint every time validation accuracy improves, so you always keep the best version.
- Run the 10-epoch fine-tuning loop by pasting this into a new cell:
best_val_acc = 0.0
# Fine-tune for 10 epochs, saving the best model
for epoch in range(10):
model.train()
correct, total = 0, 0
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
_, preds = outputs.max(1)
total += labels.size(0)
correct += preds.eq(labels).sum().item()
train_acc = 100. * correct / total
# Check validation accuracy
model.eval()
val_correct, val_total = 0, 0
with torch.no_grad():
for images, labels in val_loader:
images, labels = images.to(device), labels.to(device)
_, preds = model(images).max(1)
val_total += labels.size(0)
val_correct += preds.eq(labels).sum().item()
val_acc = 100. * val_correct / val_total
# Save when validation improves
if val_acc > best_val_acc:
best_val_acc = val_acc
torch.save(model.state_dict(), 'best_model.pth')
scheduler.step()
print(f'Epoch [{epoch+1}/10] Train: {train_acc:.2f}% | Val: {val_acc:.2f}%')
What does this training loop do?
Each epoch trains on the full augmented dataset, then evaluates on the validation split. When validation accuracy surpasses the previous best, the model weights are saved to best_model.pth.
The scheduler reduces learning rates after every 3 epochs (at epochs 3, 6, and 9), helping the model settle into a good solution as training progresses.
This cell takes a few minutes to run on a T4 GPU. You'll see epoch-by-epoch output showing both training and validation accuracy climbing.
Training seems stuck or accuracy is low?
- Verify your runtime is still using a GPU. Go to **Runtime**, then **Change runtime type**, and confirm T4 GPU is selected.
- If your session timed out, you'll need to re-run all previous cells (data loading, model setup, augmented transforms, optimizer) before this cell.
- Check that criterion is defined. It should be nn.CrossEntropyLoss() from Step 2.
Help me debug my training loop that seems stuck at low accuracy.
Before you run the evaluation, predict: do you think the Glioma/Meningioma confusion has improved, or will the same pattern persist?
- Load the best checkpoint and run inference on the test set by running this cell:
# Load the best checkpoint for final evaluation
model.load_state_dict(torch.load('best_model.pth'))
model.eval()
all_preds = []
all_labels = []
with torch.no_grad():
for images, labels in test_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
_, preds = outputs.max(1)
all_preds.extend(preds.cpu().numpy())
all_labels.extend(labels.cpu().numpy())
# Calculate overall test accuracy
accuracy = 100. * sum(p == l for p, l in zip(all_preds, all_labels)) / len(all_labels)
print(f'Test accuracy: {accuracy:.2f}%')
Why load the checkpoint first?
The training loop saved the model at its peak validation accuracy. By the end of 10 epochs, the model might have slightly overfit. Loading best_model.pth ensures you evaluate the version that generalized best.
- Generate the updated confusion matrix and classification report by running this cell:
from sklearn.metrics import confusion_matrix, classification_report
import seaborn as sns
class_names = ['Glioma', 'Meningioma', 'No Tumor', 'Pituitary']
cm = confusion_matrix(all_labels, all_preds)
# Plot confusion matrix as annotated heatmap
plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
xticklabels=class_names, yticklabels=class_names)
plt.xlabel('Predicted')
plt.ylabel('Actual')
plt.title('Confusion Matrix - Fine-Tuned Model')
plt.tight_layout()
plt.show()
# Per-class precision, recall, and F1
print(classification_report(all_labels, all_preds, target_names=class_names))
What should I see?
Compare this confusion matrix to the one from Step 3. The Glioma/Meningioma cells (the ones that had the most misclassifications) should show noticeably smaller numbers now.
Your overall test accuracy should exceed 93%. The classification report will show improved precision and recall for both Glioma and Meningioma compared to your baseline.
Your fine-tuned model now distinguishes between the previously confused tumor types with far fewer errors. The combination of augmentation and deeper fine-tuning gave the model the sensitivity it needed for subtle MRI differences.
Accuracy below 93%?
- Check that you loaded best_model.pth before running inference. Evaluating the final epoch weights (not the best checkpoint) often gives a lower number.
- Ensure your augmented transforms are only applied to the training set. The validation and test sets should still use the basic resize + normalize transforms from Step 2.
- If your runtime restarted, re-run the training loop first. The best_model.pth file is stored in the Colab filesystem and is lost on restart.
Help me figure out why my fine-tuned brain tumor model accuracy is below 93%.
✔️ Awesome, I've got everything!
Great. Double check you see 93%+ accuracy in your output and fewer Glioma/Meningioma misclassifications in the confusion matrix.
ⓧ I'd like to double check the full code
Scroll up through this step and confirm you have these cells in order:
- Augmented training transforms + recreated train_loader.
- Unfrozen layer3 and layer4 with optimizer param groups and StepLR scheduler.
- 10-epoch training loop saving best_model.pth on best validation accuracy.
- Best model loaded + test inference loop.
- Confusion matrix heatmap + classification report.
Your model now handles the toughest classification challenge in this dataset. Next up, you'll wrap this model in a web app so anyone can upload an MRI and get predictions instantly.
Deploy as a Gradio Web App
Your model achieves 93%+ accuracy on brain tumor classification. Right now, though, it lives inside a Colab notebook that only you can run.
Gradio lets you wrap any Python function into a web app with a shareable URL. In this step, you'll build an interface where anyone can upload an MRI scan and get a prediction with confidence scores, all through their browser.
In this step, get ready to:
- Write an inference function that preprocesses and classifies uploaded MRI images.
- Build a Gradio web interface displaying predictions with confidence scores.
- Launch a public URL anyone can access to try your classifier.
Install Gradio and load your model for inference
Before writing the app, you need Gradio installed and your trained model loaded into memory from the saved checkpoint.
- Install Gradio in your Colab notebook by running this command in a new cell:
!pip install gradio
Why install in Colab?
Colab comes with PyTorch pre-installed, but not Gradio. This one-line install adds the library to your current session.
- Confirm you see a success message like Successfully installed gradio in the output.
- Load your trained model from the checkpoint by running this code in a new cell:
from torchvision import models
import torch
# Load the same ResNet-50 architecture with 4-class output
model = models.resnet50()
model.fc = torch.nn.Linear(model.fc.in_features, 4)
# Load the trained weights from your best checkpoint
model.load_state_dict(torch.load('best_model.pth', map_location='cpu'))
model.eval()
print("Model loaded successfully!")
What does this code do?
This recreates the same ResNet-50 architecture you trained (with a 4-class output layer), then fills it with the weights from best_model.pth.
The map_location='cpu' ensures the model loads onto CPU, which is fast enough for classifying single images. Calling .eval() switches the model from training mode to inference mode, disabling dropout and batch norm updates.
- Confirm you see Model loaded successfully! printed in the output.
Seeing a FileNotFoundError?
This means best_model.pth is not in your current directory. Run !ls best_model.pth to check. If the file is missing, go back to Step 4 and re-run the training cell that saves the best checkpoint.
help me figure out why my model file isn't loading
Write the prediction function
Gradio calls a single Python function whenever a user uploads an image. This function needs to preprocess the image exactly like training data, run it through the model, and return a dictionary mapping class names to confidence scores.
- Define your prediction function by running this code in a new cell:
from torchvision import transforms
from PIL import Image
import os
# Class names matching the dataset folder order (alphabetical)
class_names = ['Glioma', 'Meningioma', 'No Tumor', 'Pituitary']
# Same normalization used during training (ImageNet statistics)
inference_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])
])
def predict(image):
# Preprocess the uploaded PIL image
img_tensor = inference_transform(image).unsqueeze(0)
# Run inference without computing gradients
with torch.no_grad():
outputs = model(img_tensor)
probs = torch.nn.functional.softmax(outputs, dim=1)[0]
# Return class-to-probability mapping
return {class_names[i]: float(probs[i]) for i in range(len(class_names))}
What does this code do?
The class_names list matches the alphabetical folder order that PyTorch's ImageFolder used during training. This ensures index 0 maps to Glioma, index 1 to Meningioma, and so on.
The inference_transform applies the same resize and ImageNet normalization the model was trained with. Without this, pixel values would be on the wrong scale and predictions would be random.
The predict function takes a PIL image, transforms it into a tensor, runs a forward pass, and converts raw outputs into probabilities using softmax. It returns a dictionary that Gradio can display directly.
- Test the function on one image from your test set by running this in a new cell:
# Quick test with one image from the test set
test_path = 'brain-tumor-mri/Testing/Glioma/'
test_file = os.listdir(test_path)[0]
test_image = Image.open(os.path.join(test_path, test_file))
print(f"Testing with: {test_file}")
print(predict(test_image))
- Confirm you see a dictionary printed with four class names and their probability values. The highest probability should be for Glioma since you loaded a Glioma test image.
Seeing all probabilities near 0.25?
If all four classes show roughly equal probability, the model weights may not have loaded correctly. Double check that model.eval() was called after loading the state dict. Also verify you're loading best_model.pth (not an untrained model).
help me debug why my model predictions are random
Create and launch the Gradio app
With a working prediction function, you can now wrap it in a Gradio interface. This creates a drag-and-drop upload area for images and displays the classification results as a labeled bar chart.
- Create and launch your Gradio web app by running this code in a new cell:
import gradio as gr
# Build the web interface
demo = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=4),
title="Brain Tumor MRI Classifier",
description="Upload a brain MRI scan to classify it as Glioma, Meningioma, Pituitary tumor, or No Tumor."
)
# Launch with a public shareable link (valid for 72 hours)
demo.launch(share=True)
What does this code do?
The gr.Interface connects your predict function to a web UI. It handles image upload, calls your function with the PIL image, and renders the returned dictionary as confidence bars.
Setting type="pil" on the input ensures Gradio passes a PIL Image object directly to your function. The gr.Label(num_top_classes=4) output shows all four classes ranked by confidence.
The share=True flag generates a public URL that anyone on the internet can access for 72 hours, even without a Google account or Colab access.
Before you look at the output, predict what you'll see. Will the cell print a URL, show an error, or display the app inline?
- Check the cell output. You should see two things: a Gradio app rendered inline in the notebook, and a public URL printed in the format https://xxxxx.gradio.live.
- Copy the public URL and open it in a new browser tab.
- Upload a test MRI image. You can find images in your brain-tumor-mri/Testing/ folder. Download one from Colab by right-clicking the file in the file browser sidebar and selecting Download.
- Confirm you see the predicted class name with a confidence bar chart showing all 4 classes.
- Test with at least one image from each of the 4 classes (Glioma, Meningioma, No Tumor, Pituitary) to verify the model handles all categories.
Not seeing a public URL?
If you only see a local URL (127.0.0.1) without a gradio.live link, Colab may have blocked the share tunnel. Try re-running the cell. If the issue persists, the inline notebook version still works for testing.
help me get a public Gradio share link working in Colab
✔️ Awesome, I've got everything!
Great. Double check you've saved your notebook (Ctrl+S or Cmd+S) with all four cells from this step.
ⓧ I'd like to double check the full code
Scroll up through this step and confirm you have these four cells in order, each producing output:
- Cell 1: !pip install gradio showing a successful install.
- Cell 2: Model loading code printing Model loaded successfully!.
- Cell 3: The predict function definition followed by a quick test printing a probability dictionary.
- Cell 4: The gr.Interface creation and demo.launch(share=True) showing a public URL.
Secret mission
Add Grad-CAM Explainability
Your classifier works, but can you prove WHY it made a decision? In medicine, a prediction without evidence is useless. Add Grad-CAM visualization so your app shows a heatmap of exactly where in the MRI the model is looking to make its classification.
Clean Up Your Resources
Clean Up Your Resources
Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project runs entirely on free resources (Google Colab, Kaggle, and Gradio share links), so there are no ongoing costs.
Resources you used:
- Google Colab notebook with all training and inference code (free, stored in Google Drive).
- brain-tumor-mri/ dataset directory (stored in Colab runtime, deleted when runtime disconnects).
- best_model.pth model checkpoint (stored in Colab runtime, deleted when runtime disconnects).
- Temporary Gradio share link (auto-expires after 72 hours).
- Kaggle API token uploaded to Colab (session-only, deleted when runtime disconnects).
Keep everything running
No action needed. Choose this if you plan to keep building on this project or want to push your work to GitHub and HuggingFace Spaces for a permanent portfolio piece.
- Your Colab notebook auto-saves to Google Drive. You can reopen it anytime from drive.google.com.
- The Gradio share link stays active for 72 hours. After that, restart the Gradio cell to generate a fresh link.
- Your model checkpoint and dataset persist as long as the Colab runtime is connected. Save best_model.pth to Google Drive if you want it to survive a runtime reset.
Pause - I'll come back to this later
Shut down running processes to free up resources. Your notebook and code remain safe in Google Drive.
- The Gradio share link expires automatically after 72 hours. No action needed to pause it.
- Colab will disconnect after an idle timeout (usually 90 minutes of inactivity). When it disconnects, the runtime storage is wiped, including your dataset and model file.
- To preserve your model, download best_model.pth to your local machine or copy it to Google Drive before disconnecting.
- When you return, reopen the notebook from Google Drive, reconnect the runtime, and re-run your cells.
Delete - I don't want to use this again
Remove all project resources and start fresh. No paid resources exist, so this is purely for tidiness.
- In Google Colab, click Runtime in the top menu, then click Disconnect and delete runtime. This wipes the brain-tumor-mri/ dataset, best_model.pth, and the Kaggle token from the session.
- To delete the notebook itself, go to drive.google.com. Find your notebook file (.ipynb extension, usually in a Colab Notebooks folder). Right-click it and select Remove.
- To permanently free the storage, open Trash in the Google Drive sidebar and click Empty Trash.
- The Gradio share link expires on its own after 72 hours. No manual cleanup needed.
That's a Wrap!
That's a Wrap!
Nice work! You just built a complete medical image classification pipeline, from raw MRI data through model training, evaluation, and live deployment.
You've learned how to:
- Train a ResNet-50 transfer learning model on brain MRI data, achieving over 93% accuracy across 4 tumor classes by freezing, then selectively unfreezing backbone layers.
- Diagnose model weaknesses using a confusion matrix and per-class metrics, then fix them with targeted data augmentation and deeper fine-tuning.
- Deploy a live web app using Gradio where anyone can upload an MRI scan and get instant predictions with confidence scores across all 4 classes.
- Secret Mission: Added Grad-CAM explainability to show where the model focuses in the MRI, a critical requirement for medical AI trust.
Ready to quiz yourself?