wentao yao

2021 · Classifying lung X-rays

COVID-19 detection from chest X-rays

From a KNN baseline to transfer learning with VGG16: 97.25% test accuracy in half the training time.

Year
2021
Who
Group project · my part
Stack
Deep learning · CNN · VGG16 · Transfer learning
Dataset ↗Dataset paper ↗
Chest X-rays from the dataset, labelled negative and positive.
Chest X-rays from the dataset, labelled negative and positive.

01WhatCan a model see it in the lungs?

Our group chose a project related to the pandemic: detecting COVID-19 in patients, framed as classification. This page covers my part, the chest X-ray images. The data is COVIDx CXR-2 from Kaggle: 16,490 positive and 13,992 negative training images, a balanced test set of 200 + 200, each 200 × 200 pixels.

02HowThree models, one question.

  • KNN (K = 2) as the baseline, on scaled images and without augmentation, since shifted images would distort the distances.
  • Our own CNN — two convolution layers (kernel 5), two max-pooling layers, two dense layers, ReLU and a sigmoid output, with dropout, batch normalization and augmentation.
  • VGG16 transfer learning — a pre-trained VGG16 plus two dense layers.
Scaled (left) and VGG16 preprocessing (right).
Scaled (left) and VGG16 preprocessing (right).

03ResultsTransfer learning wins twice.

0.69
KNN baseline, test accuracy
95.25%
own CNN, test — 3.6 h of training
97.25%
VGG16, test — 1.8 h of training

Transfer learning was both more accurate and twice as fast to train. One caveat: a downsampling mistake left the training data unbalanced, so all three numbers have room to improve.

VGG16 training: accuracy and loss.
VGG16 training: accuracy and loss.

04ThoughtsShadows.

Even by eye, positive X-rays often show a “shadow” in the lungs, which is why letting a model learn from the images seemed promising — and 97.25% suggests it is feasible. If it really is about spotting shadows, an object detector like YOLOv5 could be the next thing to try.

Adapted from my original write-up · more in the links above.