In early 2022, during my final year of undergrad at Aligarh Muslim University, COVID-19 was still overwhelming healthcare systems in many parts of the world. The gold-standard diagnostic test — RT-PCR — was accurate but slow (results took 24-48 hours) and required specialized lab equipment that wasn't available in rural clinics. Chest X-rays, however, were fast, cheap, and available almost everywhere.
The question was simple: could a deep learning model read a chest X-ray and reliably detect COVID-19 — fast enough to aid triage decisions, and lightweight enough to run on a hospital's existing hardware?
That question became my undergraduate research project, led to two published papers, and taught me more about real-world deep learning than any textbook could. This is the full story.
1. The Problem
COVID-19 produces characteristic patterns in chest X-rays — bilateral ground-glass opacities, peripheral consolidation, and in severe cases, a "white-out" appearance. Experienced radiologists can spot these patterns, but the volume of cases during peak waves far exceeded available radiologist capacity.
We set out to build a classification system that could categorize CXR images into three classes:
- COVID-19 positive — characteristic COVID pneumonia patterns
- Non-COVID pneumonia — bacterial or other viral pneumonia (important to distinguish from COVID)
- Normal — healthy lungs with no pathological findings
The three-class problem is harder than simple binary (COVID vs. normal) because the model needs to distinguish between two types of pneumonia that can look similar to the untrained eye. But it's clinically more useful — a doctor needs to know not just "is something wrong?" but "is it COVID specifically?"
2. Dataset
Data is the hardest part of any medical AI project. We assembled our dataset from multiple public sources to ensure diversity and reduce institutional bias:
- COVID-19 Radiography Database (Qatar University & University of Dhaka) — 3,616 COVID-positive CXR images
- RSNA Pneumonia Detection Challenge — non-COVID pneumonia and normal cases
- NIH ChestX-ray14 — additional normal cases for balance
After curation and quality checks (removing duplicates, discarding low-quality scans, ensuring consistent image orientation), our final dataset contained:
# Dataset distribution
| Class | Train | Validation | Test | Total |
|--------------------|--------|------------|-------|--------|
| COVID-19 | 2,892 | 362 | 362 | 3,616 |
| Non-COVID Pneumonia| 2,900 | 363 | 363 | 3,626 |
| Normal | 2,908 | 364 | 364 | 3,636 |
|--------------------|--------|------------|-------|--------|
| Total | 8,700 | 1,089 | 1,089 | 10,878 |We deliberately balanced the classes to prevent the model from developing a bias toward the majority class. All images were resized to 224x224 pixels and normalized to the ImageNet mean and standard deviation.
Data augmentation
Medical imaging datasets are small by deep learning standards. To prevent overfitting and improve generalization, we applied aggressive but clinically valid augmentations:
import albumentations as A
train_transforms = A.Compose([
A.RandomResizedCrop(224, 224, scale=(0.85, 1.0)),
A.HorizontalFlip(p=0.5),
A.RandomBrightnessContrast(
brightness_limit=0.15,
contrast_limit=0.15, p=0.5
),
A.GaussNoise(var_limit=(5, 25), p=0.3),
A.Rotate(limit=10, p=0.5),
A.GaussianBlur(blur_limit=3, p=0.2),
A.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
),
ToTensorV2()
])A critical decision: we did not apply vertical flips or large rotations. A chest X-ray has a fixed anatomical orientation — the heart is on the left, the diaphragm is at the bottom. Flipping vertically or rotating 90 degrees would create anatomically impossible images that could confuse the model.
3. Architecture Design
Rather than using a single pre-trained model, we took a comparative approach — systematically evaluating multiple architectures to understand which design choices matter most for CXR classification. This comparative analysis formed the basis of our first paper (Taylor & Francis, 2023).
Models evaluated
- VGG-16 — deep but computationally expensive, lots of parameters
- ResNet-50 — skip connections for deeper training without degradation
- DenseNet-121 — feature reuse through dense connections; popular in medical imaging
- EfficientNet-B0 — compound scaling for balanced depth/width/resolution
- MobileNetV2 — lightweight depthwise separable convolutions; designed for mobile
- Custom CNN — our own architecture optimized for this specific task
All transfer learning models were initialized with ImageNet weights and fine-tuned end-to-end. We replaced the final classification head with a custom head: Global Average Pooling, Dropout(0.4), Dense(256, ReLU), Dropout(0.3), Dense(3, Softmax).
Our custom CNN
The custom architecture was designed with two goals: maximize accuracy and minimize model size for edge deployment. We used a modular block design inspired by EfficientNet's compound scaling but tailored to the CXR domain:
import torch.nn as nn
class CXRBlock(nn.Module):
"""Residual block with depthwise separable convolutions."""
def __init__(self, in_ch, out_ch, stride=1):
super().__init__()
self.conv = nn.Sequential(
# Depthwise
nn.Conv2d(in_ch, in_ch, 3, stride, 1, groups=in_ch, bias=False),
nn.BatchNorm2d(in_ch),
nn.ReLU6(inplace=True),
# Pointwise
nn.Conv2d(in_ch, out_ch, 1, bias=False),
nn.BatchNorm2d(out_ch),
)
self.skip = (
nn.Sequential(nn.Conv2d(in_ch, out_ch, 1, stride, bias=False),
nn.BatchNorm2d(out_ch))
if in_ch != out_ch or stride != 1 else nn.Identity()
)
self.relu = nn.ReLU6(inplace=True)
def forward(self, x):
return self.relu(self.conv(x) + self.skip(x))
class COVID_CXR_Net(nn.Module):
def __init__(self, num_classes=3):
super().__init__()
self.features = nn.Sequential(
# Stem
nn.Conv2d(3, 32, 3, 2, 1, bias=False),
nn.BatchNorm2d(32),
nn.ReLU6(inplace=True),
# Blocks with progressive channel expansion
CXRBlock(32, 64, stride=2),
CXRBlock(64, 64),
CXRBlock(64, 128, stride=2),
CXRBlock(128, 128),
CXRBlock(128, 256, stride=2),
CXRBlock(256, 256),
CXRBlock(256, 512, stride=2),
)
self.classifier = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Dropout(0.4),
nn.Linear(512, 256),
nn.ReLU6(inplace=True),
nn.Dropout(0.3),
nn.Linear(256, num_classes)
)
def forward(self, x):
return self.classifier(self.features(x))Key design choices: depthwise separable convolutions reduce parameters by ~8-9x compared to standard convolutions while maintaining representational power. ReLU6 bounds activations for better quantization later. Residual connections in every block prevent gradient degradation.
4. Training
We trained all models using the same protocol for fair comparison:
# Training configuration
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
optimizer, T_0=10, T_mult=2
)
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
# Training: 50 epochs, early stopping with patience=7
# Hardware: NVIDIA Tesla T4 (Google Colab Pro)Label smoothing (0.1) was important — it prevents the model from becoming overconfident on training examples, which improved generalization on our relatively small dataset. Cosine annealing with warm restarts helped escape local minima and consistently found better solutions than step-based schedulers.
Class activation maps
To verify the model was learning clinically meaningful features (not just artifacts), we generated Grad-CAM visualizations. These heatmaps show which regions of the X-ray the model focuses on when making its prediction:
- For COVID-positive cases, the model consistently highlighted bilateral peripheral regions — exactly where ground-glass opacities appear
- For bacterial pneumonia, activation focused on unilateral lower lobes — matching the typical lobar consolidation pattern
- For normal cases, activations were diffuse and low-intensity — the model correctly found "nothing to see"
This was a crucial validation step. In medical imaging, a model that gets the right answer for the wrong reason is dangerous. If the model had been keying on image artifacts (like text overlays or equipment markers), the Grad-CAM maps would have revealed that immediately.
5. Results — Comparative Analysis
Here's how each architecture performed on our held-out test set of 1,089 images:
# Test set performance (1,089 images)
| Model | Accuracy | Precision | Recall | F1-Score | Params |
|----------------|----------|-----------|--------|----------|----------|
| Custom CNN | 97.2% | 97.3% | 97.1% | 97.2% | 2.1M |
| EfficientNet-B0| 96.8% | 96.9% | 96.7% | 96.8% | 5.3M |
| DenseNet-121 | 96.4% | 96.5% | 96.3% | 96.4% | 7.0M |
| ResNet-50 | 95.9% | 96.1% | 95.8% | 95.9% | 23.5M |
| MobileNetV2 | 95.1% | 95.3% | 94.9% | 95.1% | 2.2M |
| VGG-16 | 93.7% | 93.9% | 93.5% | 93.7% | 134.3M |The standout result: our custom CNN achieved the highest accuracy (97.2%) with only 2.1M parameters — 3x smaller than EfficientNet-B0 and 64x smaller than VGG-16.
Per-class performance
# Custom CNN — Per-class metrics
| Class | Precision | Recall | F1-Score |
|---------------------|-----------|--------|----------|
| COVID-19 | 98.1% | 97.5% | 97.8% |
| Non-COVID Pneumonia | 96.2% | 96.8% | 96.5% |
| Normal | 97.5% | 97.0% | 97.2% |The hardest distinction was between COVID and non-COVID pneumonia (as expected), but even there we achieved >96% precision and recall. The confusion matrix showed that the few misclassifications were almost exclusively between the two pneumonia classes — the model never confused pneumonia with normal lungs.
6. Model Compression for Edge Deployment
A model that only runs on a GPU server isn't useful in a rural clinic. Our second paper (ICBSII 2024) focused on making the model deployable on edge devices — think Raspberry Pi, Jetson Nano, or even a smartphone.
We applied three compression techniques sequentially:
Step 1: Pruning
We used structured pruning to remove entire filters (channels) that contribute least to the output:
import torch.nn.utils.prune as prune
# Prune 30% of channels by L1 norm
for name, module in model.named_modules():
if isinstance(module, nn.Conv2d):
prune.ln_structured(
module, name='weight',
amount=0.3, n=1, dim=0
)
# Fine-tune to recover accuracy
# 5 epochs, lower learning rate 1e-4After pruning 30% of filters and fine-tuning, accuracy dropped only 0.4% (97.2% to 96.8%), but the model was 40% faster at inference.
Step 2: Quantization
We applied post-training quantization to convert 32-bit floating point weights to 8-bit integers — cutting model size by ~4x:
import torch.quantization as quant
# Prepare for quantization
model.eval()
model.qconfig = quant.get_default_qconfig('fbgemm')
model_prepared = quant.prepare(model)
# Calibrate with representative data (100 batches)
with torch.no_grad():
for batch in calibration_loader:
model_prepared(batch)
# Convert to quantized model
model_quantized = quant.convert(model_prepared)
# Size: 8.4 MB (FP32) -> 2.2 MB (INT8)
# Compression ratio: 3.8xStep 3: Knowledge distillation
Finally, we trained a tiny student model (only 0.5M parameters) using knowledge distillation — the full custom CNN acted as the teacher:
def distillation_loss(student_logits, teacher_logits, labels,
temperature=4.0, alpha=0.7):
"""Combine soft teacher targets with hard ground truth."""
soft_loss = nn.KLDivLoss(reduction='batchmean')(
F.log_softmax(student_logits / temperature, dim=1),
F.softmax(teacher_logits / temperature, dim=1)
) * (temperature ** 2)
hard_loss = nn.CrossEntropyLoss()(student_logits, labels)
return alpha * soft_loss + (1 - alpha) * hard_lossEdge deployment results
# Compression pipeline results
| Model Variant | Accuracy | Size | Inference (RPi 4) |
|---------------------|----------|--------|--------------------|
| Full model (FP32) | 97.2% | 8.4 MB | 1.2 sec |
| Pruned + FT | 96.8% | 5.1 MB | 0.7 sec |
| Pruned + Quantized | 96.5% | 2.2 MB | 0.3 sec |
| Distilled student | 95.1% | 0.6 MB | 0.1 sec |The quantized pruned model hit the sweet spot: 96.5% accuracy in 0.3 seconds on a Raspberry Pi 4, with a model size of just 2.2 MB. That's small enough to deploy on virtually any hardware, fast enough for real-time triage, and accurate enough to be clinically useful as a screening aid.
7. Conference Presentations
This work resulted in two published papers:
Paper 1 (2023): "Comparative Analysis of Deep Learning Techniques for Fast Detection of COVID-19 Using CXR Images" — presented at the International Conference on Advances in Computational Intelligence and its Applications, published by Taylor and Francis (CRC Press). This paper covered the full comparative analysis across six architectures.
Paper 2 (2024): "Deep Learning Model for Edge Devices for COVID-19 Detection from CXR Images" — presented at the 10th International Conference on Bio Signals, Images, and Instrumentation (ICBSII) in Chennai, India. This paper focused on the compression pipeline and edge deployment.
Presenting at ICBSII in Chennai was a defining moment in my academic journey. The questions from the audience — especially from radiologists in the room — sharpened my understanding of the gap between "model accuracy" and "clinical utility." A 97% accuracy number means nothing if the doctor doesn't trust the system.
8. Reflections & What I'd Do Differently
Use Vision Transformers. When I started this project in 2022, ViTs were still emerging in medical imaging. Today, models like DeiT and Swin Transformer consistently outperform CNNs on medical classification tasks. If I were starting over, I'd benchmark a ViT-Small alongside the CNN architectures.
External validation is everything. Our test set came from the same data sources as training (different split, but same distribution). A truly robust evaluation would test on CXR images from a completely different hospital system — different X-ray machines, different patient demographics, different image quality. Distribution shift is the silent killer of medical AI models.
Uncertainty estimation matters. In a clinical setting, the model should say "I'm not sure" rather than confidently predict the wrong class. Adding MC Dropout or an ensemble for uncertainty quantification would make the system much safer for real-world triage.
ONNX for deployment. We exported the final model to TorchScript for edge deployment, but ONNX Runtime would have been a better choice — it runs on more hardware targets and generally delivers faster inference, especially with INT8 quantization on ARM processors.
What This Project Taught Me
This wasn't just an academic exercise. It was the project that convinced me to pursue graduate studies in ML and shaped how I think about building AI systems:
- The model is 20% of the work. Data curation, validation, compression, deployment, and clinical trust are the other 80%
- Small, purpose-built models can beat large general ones when you understand your domain deeply
- Always verify what the model is learning — Grad-CAM saved us from deploying a model that was cheating by reading text labels on images
- Edge deployment is a first-class requirement, not an afterthought — if your model can't run where it's needed, it doesn't matter how accurate it is
This project laid the foundation for everything I've done since — from healthcare AI (Lab Lens) to deploying production LLMs at Checkit Analytics. The core lesson remains the same: build AI that works where it's needed, not just where it's convenient.
If you're working on medical imaging or edge ML, I'd love to connect. Reach out via email or find me on GitHub.