Back to Articles AI Applications • 15 min read

Building a Medical Image Classifier: What Happened When I Tried AI on Chest X-Rays

Last year, I decided to test how easy it was to train a deep learning vision model on actual medical diagnostic data. I downloaded the public NIH Chest X-ray dataset—a massive collection of over 100,000 anonymized frontal-view X-ray images labeled with conditions like Pneumonia, Cardiomegaly, and Effusion. On paper, my plan was straightforward: fine-tune a pre-trained ResNet-50 network to detect Pneumonia.

Within two days of training my first model, I achieved a staggering 94% accuracy score. I was thrilled—until I looked deeper into the confusion matrix and discovered a terrifying flaw that taught me more about real-world AI engineering than any textbook ever could.

---

1. The 94% Accuracy Illusion

When I inspected my dataset distribution, I realized why my model was scoring 94% accuracy while being completely useless:

Out of 100,000 images, roughly 94,000 were labeled 'No Finding' (healthy lungs), and only 6,000 showed clear signs of Pneumonia. My model had learned a lazy trick: predict 'Healthy' for every single patient. It scored 94% accuracy without ever detecting a single sick patient! In medical machine learning, accuracy is a dangerously misleading metric.

---

2. The Fix: Weighted Loss & PyTorch Implementation

To force the network to care about the rare positive cases, I had to replace standard cross-entropy loss with a **Weighted Binary Cross-Entropy Loss** and track Sensitivity (Recall) rather than raw accuracy. Here is the PyTorch pipeline I wrote to fix the class imbalance:

import torch
import torch.nn as nn
from torchvision import models, transforms
from torch.utils.data import DataLoader

# Compute pos_weight based on class distribution
# Healthy = 94,000, Pneumonia = 6,000 -> weight ratio = 94000 / 6000 ≈ 15.6
pos_weight = torch.tensor([15.6]).cuda()

# Use weighted loss so missing a sick patient penalties the loss 15x harder
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)

def train_one_epoch(model, dataloader, optimizer):
    model.train()
    total_loss = 0.0
    true_positives = 0
    false_negatives = 0
    
    for images, labels in dataloader:
        images, labels = images.cuda(), labels.float().cuda()
        optimizer.zero_grad()
        
        outputs = model(images).squeeze()
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        preds = (torch.sigmoid(outputs) > 0.5).float()
        true_positives += ((preds == 1) & (labels == 1)).sum().item()
        false_negatives += ((preds == 0) & (labels == 1)).sum().item()
        
    sensitivity = true_positives / (true_positives + false_negatives + 1e-6)
    print(f"Epoch Sensitivity (Pneumonia Recall): {sensitivity:.2%}")

# After adding pos_weight, Sensitivity jumped from 0.0% to 88.4%!
            

---

3. The Shortcut Learning Shock

Even after fixing the loss function, I discovered another bizarre issue when I generated Grad-CAM saliency heatmaps to see *which pixels* the neural network was looking at to make its decision.

Instead of focusing on lung tissue anomalies, the neural network heatmaps were highlighting the **top-right corner of the X-ray images**! Why? Because portable X-ray machines used for emergency room patients had a metal 'PORTABLE' tag burned into the corner of the film. Patients who were too sick to stand up got portable X-rays. The AI was not diagnosing pneumonia—it was simply detecting the word 'PORTABLE' on the film!

---

Why Deploying Medical AI Terrifies Me

Pitfall	What Happened in My Model	How It Must Be Solved
Class Imbalance	Model predicted 'Healthy' for everyone to get 94% accuracy.	Use weighted BCE loss & ROC-AUC / Recall metrics.
Shortcut Learning	Model learned hospital tags instead of lung opacities.	Crop metadata tags & audit pixel heatmaps with Grad-CAM.
Domain Shift	Accuracy plummeted when tested on X-rays from a different hospital.	Validate across diverse external multi-center datasets.

Building a high-scoring medical AI model in PyTorch is easy. Building one that is genuinely trustworthy, un-biased, and safe for human patient care is one of the hardest engineering challenges in computer science.