Building a Pneumonia Diagnostic Assistant with PyTorch Lightning

A Deep Learning Pipeline for Medical Imaging

Introduction:

Why This Project Matters

Most machine learning projects stop at “it works on my dataset.”

Medical AI doesn’t get that luxury.

In healthcare, structure, reproducibility, and robustness matter as much as accuracy. A model that predicts pneumonia incorrectly isn’t just a bad metric — it’s a potential clinical risk.

In this project, we build the core deep learning engine of a Pneumonia Diagnostic Assistant using PyTorch Lightning, applying professional ML engineering practices from data ingestion to training orchestration.

You can find my entire code here!

The Goal

Train a deep learning model that can classify chest X-ray images into:

• Normal
• Bacterial Pneumonia
• Viral Pneumonia

This is a multiclass medical imaging problem, built using:

• Transfer learning
• Clean data abstraction
• Mixed-precision training
• Early stopping for safety

Why PyTorch Lightning?

PyTorch is powerful — but raw PyTorch code gets messy fast.

Lightning solves three real engineering problems:

1. Separation of concerns (data, model, training loop)
2. Reproducibility
3. Scalable training without boilerplate

If you’re building ML systems meant to survive beyond a notebook, Lightning is the right abstraction.

Dataset Overview: Chest X-Ray Images (Pneumonia)

You Can Get It From HERE:

[embed]Chest X-Ray Images (Pneumonia) 5,863 images, 2 categorieswww.kaggle.com

The dataset consists of pediatric chest X-ray images reviewed by medical professionals. For this assignment:

• The original PNEUMONIA class was split into:
• BACTERIAL_PNEUMONIA
• VIRAL_PNEUMONIA
• Validation and test sets were merged
• All classes were balanced to avoid bias

JUST IN CASE YOU’RE DOWNLOADING IT FROM KAGGLE YOU’LL NEED TO PREPROCESS THE DATASET.

HERE’S HOW YOU CAN DO IT

Preprocessing for Kaggle Chest X-Ray Pneumonia Dataset

What This Code Does (High-Level)

Starting from the raw Kaggle dataset, this pipeline will:

1. Load the original dataset structure
2. Split PNEUMONIA into:

• BACTERIAL_PNEUMONIA
• VIRAL_PNEUMONIA

1. Merge test + val into a single validation set
2. Balance all classes equally (train + val)
3. Produce a Lightning-ready directory structure

Expected Kaggle Dataset Structure (Raw)

After downloading from Kaggle:

chest_xray/
├── train/
│   ├── NORMAL/
│   └── PNEUMONIA/
├── test/
│   ├── NORMAL/
│   └── PNEUMONIA/
└── val/
    ├── NORMAL/
    └── PNEUMONIA/

Inside PNEUMONIA/, filenames look like:

• person23_bacteria_45.jpeg
• person89_virus_120.jpeg

That filename is how we split bacterial vs viral.

Step 0: Imports & Setup

import os
import shutil
import random
from pathlib import Path
from collections import defaultdict

Step 1: Define Paths

RAW_DATA_DIR = Path("chest_xray")          # Kaggle dataset root
OUTPUT_DATA_DIR = Path("chest_xray_clean") # Final processed dataset

TRAIN_DIR = OUTPUT_DATA_DIR / "train"
VAL_DIR = OUTPUT_DATA_DIR / "val"

CLASSES = ["NORMAL", "BACTERIAL_PNEUMONIA", "VIRAL_PNEUMONIA"]

Step 2: Utility Functions

Detect Pneumonia Type from Filename

def pneumonia_type_from_filename(filename):
    filename = filename.lower()
    if "bacteria" in filename:
        return "BACTERIAL_PNEUMONIA"
    elif "virus" in filename:
        return "VIRAL_PNEUMONIA"
    else:
        raise ValueError(f"Cannot determine pneumonia type: {filename}")

Create Output Directories

def create_output_dirs():
    for split in ["train", "val"]:
        for cls in CLASSES:
            os.makedirs(OUTPUT_DATA_DIR / split / cls, exist_ok=True)

Step 3: Collect All Images (Train + Test + Val)

We first pool everything together so we can rebalance cleanly.

def collect_images():
    collected = defaultdict(list)

    for split in ["train", "test", "val"]:
        split_path = RAW_DATA_DIR / split

        for cls in ["NORMAL", "PNEUMONIA"]:
            class_path = split_path / cls
            if not class_path.exists():
                continue

            for img in class_path.iterdir():
                if cls == "NORMAL":
                    collected["NORMAL"].append(img)
                else:
                    pneu_type = pneumonia_type_from_filename(img.name)
                    collected[pneu_type].append(img)

    return collected

Step 4: Balance the Dataset

This is critical for medical fairness.

def balance_classes(collected):
    min_count = min(len(v) for v in collected.values())
    balanced = {}

    for cls, images in collected.items():
        random.shuffle(images)
        balanced[cls] = images[:min_count]

    return balanced

Step 5: Train / Validation Split

We create a clean split AFTER balancing.

def split_train_val(balanced, val_ratio=0.2):
    train_split = {}
    val_split = {}

    for cls, images in balanced.items():
        split_idx = int(len(images) * (1 - val_ratio))
        train_split[cls] = images[:split_idx]
        val_split[cls] = images[split_idx:]

    return train_split, val_split

Step 6: Copy Files to Final Structure

def copy_files(split_dict, split_name):
    for cls, images in split_dict.items():
        for img_path in images:
            dst = OUTPUT_DATA_DIR / split_name / cls / img_path.name
            shutil.copy(img_path, dst)

Step 7: Run the Full Pipeline

def preprocess_kaggle_chest_xray():
    print("Creating output directories...")
    create_output_dirs()

    print("Collecting images...")
    collected = collect_images()

    print("Balancing classes...")
    balanced = balance_classes(collected)

    print("Splitting train / val...")
    train_split, val_split = split_train_val(balanced)

    print("Copying training files...")
    copy_files(train_split, "train")

    print("Copying validation files...")
    copy_files(val_split, "val")

    print("\n✅ Preprocessing complete!")
    for cls in CLASSES:
        print(
            f"{cls}: "
            f"train={len(train_split[cls])}, "
            f"val={len(val_split[cls])}"
        )

Step 8: Execute

preprocess_kaggle_chest_xray()

Final Output Structure (Lightning-Ready)

chest_xray_clean/
├── train/
│   ├── NORMAL/
│   ├── BACTERIAL_PNEUMONIA/
│   └── VIRAL_PNEUMONIA/
└── val/
    ├── NORMAL/
    ├── BACTERIAL_PNEUMONIA/
    └── VIRAL_PNEUMONIA/

This now plugs directly into the ChestXRayDataModule you already built.

Directory Structure

chest_xray/
├── train/
│   ├── NORMAL/
│   ├── BACTERIAL_PNEUMONIA/
│   └── VIRAL_PNEUMONIA/
└── val/
    ├── NORMAL/
    ├── BACTERIAL_PNEUMONIA/
    └── VIRAL_PNEUMONIA/

Step 1: Data Engineering with

LightningDataModule

Medical AI pipelines live or die by data hygiene.

We encapsulate all data logic inside a LightningDataModule:

• Paths
• Transforms
• Dataset creation
• DataLoaders

Image Transformations

We apply augmentation only during training, never validation — a critical medical ML rule.

TRAIN_TRANSFORM = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=(0.9, 1.1)),
    transforms.ToTensor(),
    transforms.Normalize([0.482]*3, [0.222]*3)
])

VAL_TRANSFORM = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.482]*3, [0.222]*3)
])

DataModule Implementation

class ChestXRayDataModule(pl.LightningDataModule):
    def __init__(self, data_dir, batch_size=64):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.train_transform = TRAIN_TRANSFORM
        self.val_transform = VAL_TRANSFORM

    def setup(self, stage=None):
        self.train_dataset = datasets.ImageFolder(
            os.path.join(self.data_dir, "train"),
            transform=self.train_transform
        )
        self.val_dataset = datasets.ImageFolder(
            os.path.join(self.data_dir, "val"),
            transform=self.val_transform
        )

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=self.batch_size)

Why this matters:

You can now swap datasets, batch sizes, or transforms without touching model code.

Step 2: Model Architecture with Transfer Learning

Training a CNN from scratch on medical data is:

• Slow
• Data-hungry
• Often unnecessary

Instead, we use ResNet-18 as a feature extractor.

Key Strategy

• Load pretrained weights
• Freeze all layers
• Train only the classifier head

This minimizes overfitting and speeds up convergence.

def load_resnet18(num_classes, weights_path):
    model = tv_models.resnet18(weights=None)
    model.fc = nn.Linear(model.fc.in_features, num_classes)

    state_dict = torch.load(weights_path, map_location="cpu")
    model.load_state_dict(state_dict)

    for p in model.parameters():
        p.requires_grad = False
    for p in model.fc.parameters():
        p.requires_grad = True

    return model

Step 3: The

LightningModule

— Model Brain

The LightningModule defines:

• Forward pass
• Loss
• Metrics
• Optimizers
• Schedulers

Why This Design Is Professional

Everything related to learning lives in one place.

No training logic leaks into notebooks or scripts.

Implementation

class ChestXRayClassifier(pl.LightningModule):
    def __init__(self, model_weights_path, num_classes=3, learning_rate=1e-3, weight_decay=1e-2):
        super().__init__()
        self.save_hyperparameters()

        self.model = load_resnet18(num_classes, model_weights_path)
        self.loss_fn = nn.CrossEntropyLoss()
        self.accuracy = Accuracy(task="multiclass", num_classes=num_classes)

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = self.loss_fn(logits, y)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = self.loss_fn(logits, y)
        acc = self.accuracy(logits, y)
        self.log_dict({"val_loss": loss, "val_acc": acc}, prog_bar=True)

    def configure_optimizers(self):
        optimizer = optim.AdamW(self.parameters(), lr=self.hparams.learning_rate,
                                weight_decay=self.hparams.weight_decay)
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode="min", factor=0.1, patience=2
        )
        return {"optimizer": optimizer,
                "lr_scheduler": {"scheduler": scheduler, "monitor": "val_loss"}}

Step 4: Early Stopping — Safety First

Overtraining medical models is dangerous.

We use EarlyStopping based on validation accuracy.

def early_stopping(num_epochs, stop_threshold):
    return EarlyStopping(
        monitor="val_acc",
        stopping_threshold=stop_threshold,
        patience=int(num_epochs / 2),
        mode="max"
    )

This ensures:

• Training stops once performance is “good enough”
• No wasted compute
• Reduced overfitting risk

Step 5: Training Orchestration with Mixed Precision

Modern GPUs are built for mixed precision. Not using it is leaving performance on the table.

Trainer Configuration

def run_training(model, data_module, num_epochs, callback, progress_bar=True, dry_run=False):
    trainer = pl.Trainer(
        max_epochs=num_epochs,
        accelerator="auto",
        devices=1,
        precision="16-mixed",
        callbacks=[callback],
        logger=False,
        enable_progress_bar=progress_bar,
        enable_model_summary=False,
        enable_checkpointing=False,
        fast_dev_run=dry_run
    )

    trainer.fit(model, data_module)
    return trainer, model

Why This Setup Works

• 16-bit precision → faster + less memory
• Dry runs → sanity checks
• No hidden side effects (logging/checkpoints disabled)

Final Thoughts: From Notebook to Medical AI System

This project isn’t about squeezing out another 0.5% accuracy.

It’s about:

• Building maintainable ML systems
• Respecting medical constraints
• Writing code that another engineer (or regulator) can understand

You now have:

• A clean data pipeline
• A transfer-learning classifier
• Safe training controls
• A scalable Lightning architecture

This is the foundation that you can use in real medical AI teams.