XR-AI: Chest X-Ray Disease Classification

Transfer learning model for automated chest X-ray classification. Identifies tuberculosis, pneumonia, COVID-19, and normal cases using DenseNet-121 pretrained on ImageNet with Grad-CAM visualization for interpretability.

Built in 48 hours at KU Hackfest 2025 by Abhyudaya Pokhrel, Unique Shrestha, and myself.

Overview

XR-AI applies transfer learning to the problem of chest X-ray classification. Medical imaging datasets are expensive to label and often small. Transfer learning addresses this by adapting a model already trained on millions of images (ImageNet) to the target domain, requiring far less labeled medical data to achieve strong performance.

The model classifies chest X-rays into four categories:

• Tuberculosis
• Pneumonia
• COVID-19
• Normal

Achieved 90% accuracy on chest X-ray test datasets.

Architecture

Final architecture: DenseNet-121

DenseNet (Densely Connected Convolutional Network) connects each layer to every subsequent layer in the network. Where ResNet uses additive skip connections between individual layers, DenseNet concatenates feature maps from all preceding layers into each subsequent layer. This means early features (low-level edges and textures) remain directly accessible to deep layers.

For medical imaging this is particularly useful. Disease signatures in chest X-rays often manifest as subtle texture and structural changes that benefit from multi-scale feature availability throughout the network.

Input X-ray (224×224×3)
    ↓
DenseNet-121 backbone (pretrained ImageNet weights, frozen initially)
    ↓
Global Average Pooling
    ↓
Custom classifier head
    ↓
FC(1024 → 256) → ReLU → Dropout(0.5)
    ↓
FC(256 → 4) → Softmax
    ↓
Output: [TB, Pneumonia, COVID-19, Normal]

Architecture Decision: ResNet vs DenseNet

We started with ResNet as a baseline. ResNet's residual connections address vanishing gradients effectively and it performs well on general image classification. However, we observed:

• Lower validation accuracy overall
• Weak class separation between COVID-19 and pneumonia
• Less stable training curves

Transitioning to DenseNet-121 produced noticeable improvement in validation performance and training stability. The dense feature reuse appears better suited to the fine-grained visual differences between COVID-19 and pneumonia presentations.

Training Details

Normalization (ImageNet statistics):

mean = [0.485, 0.456, 0.406]
std  = [0.229, 0.224, 0.225]

Proper normalization to ImageNet statistics was significant for training stability. The pretrained weights learned features from data normalized this way; mismatched normalization degrades transfer.

Data augmentation:

transforms.RandomHorizontalFlip()
transforms.RandomRotation(10)
transforms.ColorJitter(brightness=0.2, contrast=0.2)

Grad-CAM Visualization

Gradient-weighted Class Activation Mapping (Grad-CAM) highlights which spatial regions of an input image most influenced the model's classification decision.

For a given class prediction, Grad-CAM computes the gradient of the class score with respect to the final convolutional feature maps. These gradients are globally averaged to produce weights, which are then used to produce a weighted combination of the feature maps. The result is a coarse localization heatmap overlaid on the original image.

This interpretability is important in clinical contexts: a model that correctly classifies for incorrect anatomical reasons cannot be trusted in deployment.

Observed heatmap patterns:

• Tuberculosis: concentration in upper lung lobes (consistent with typical TB pathology)
• Pneumonia: diffuse focus in lower lobes
• COVID-19: bilateral peripheral patterns

References

• Huang et al. (2017). Densely Connected Convolutional Networks. CVPR.
• Selvaraju et al. (2017). Grad-CAM: Visual Explanations from Deep Networks. ICCV.
• He et al. (2016). Deep Residual Learning for Image Recognition. CVPR.

Team

• Abhyudaya Pokhrel
• Unique Shrestha

Built at KU Hackfest 2025 (48 hours).