Pneumothorax Classification

Computer Vision Engineer

Impact

Automated detection of collapsed lungs in Chest X-Rays using EfficientNet-B4.

Overview

The Problem: Critical Seconds in Radiology

Pneumothorax—a collapsed lung—is a life-threatening condition that requires immediate intervention. In busy emergency rooms, radiologists are under immense pressure, and subtle signs in Chest X-rays can be missed. A delay in diagnosis can be fatal. We needed an automated system to act as a "second opinion," instantly flagging high-risk scans for priority review.

The Solution: Deep Learning at the Edge

We developed a high-performant binary classification model capable of distinguishing between healthy lungs and those with Pneumothorax. By leveraging Transfer Learning, we adapted state-of-the-art vision models to the medical domain, achieving robust accuracy with limited labeled data.

Technical Deep Dive

1. Model Architecture: EfficientNet-B4

We chose EfficientNet-B4 for its superior balance of accuracy and efficiency.

Figure: EfficientNet Architecture

Pre-training: Trained on ImageNet-1k at a resolution of 380x380.
Compound Scaling: Unlike traditional ConvNets that scale depth, width, or resolution arbitrarily, EfficientNet proposes a compound coefficient to uniformly scale all dimensions. This results in a "mobile-friendly" pure convolutional model that punches far above its weight class.
Backbone: timm/tf_efficientnet_b4.ns_jft_in1k

2. Custom Classification Head

The original ImageNet classifier isn't suitable for binary medical diagnosis. We engineered a custom head:

Dropout(0.25): For regularization.
Linear projection: Maps features to 512 dimensions.
BatchNorm1d & ReLU
Final Linear layer: Outputs the probability of Pneumothorax presence.

3. The Training Engine: PyTorch Lightning

We strictly adhered to MLOps principles using PyTorch Lightning to decouple engineering from research code.

Mixed Precision (16-mixed): Accelerated training speed and reduced memory usage without sacrificing convergence.
Smart Scheduling: Implemented ReduceLROnPlateau to dynamically adjust the learning rate when validation loss plateaus.
Early Stopping: Automatically halts training when the model stops improving, preventing overfitting.

4. Evaluation Metrics

In medical diagnosis, accuracy isn't enough. We optimized for a comprehensive metric suite:

BinaryF1Score: To balance precision and recall.
BinaryPrecision & Recall: ensuring we minimize false negatives (critical errors).

Stack

Framework: PyTorch, PyTorch Lightning
Models: timm (EfficientNet), torchvision
Optimization: AdamW, OneCycleLR
Logging: TensorBoard

Computer VisionPyTorch LightningMedical AIEfficientNet

Gallery Overview

Pneumothorax Classification gallery image

View Project Site