COVID-19 Detection from Chest X-rays Using Transfer Learning

Problem Statement

Rapid detection of COVID-19 cases is critical to containing the virus as well as decreasing pressure on an overflowing healthcare system. Detection of COVID-19 has a lengthy clinical testing time, thus imaging tools, such as Chest X-rays, are a key instrument in helping to detect COVID-19 and speeding up the identification process. In this notebook, we train a deep convolutional neural network using transfer learning on a novel dataset of 15,000 chest X-rays to aid in the detection of COVID-19. Further, we implement local interpretable model-agnostic explanations (LIME) to provide insight into the model's predictions and enabling healthcare providers to make both time-efficient and correct diagnosis.

'''Load PyTorch and necessary libraries.'''

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torch.nn.functional as F

import torchvision
import torchvision.datasets as datasets
from torchvision import models
import torchvision.transforms as transforms

import os
from pathlib import Path
import Augmentor
from IPython.display import Image, display
import splitfolders
import matplotlib.pyplot as plt
import numpy as np
import time
from sklearn.metrics import f1_score, precision_score, recall_score, classification_report
import seaborn as sns
import pandas as pd
from PIL import Image
dir = os.getcwd()

Data Collection and Augmentation

We use chest X-ray images from 3 categories: normal, pneumonia, and COVID-19 cases. The dataset, containing 5000 normal images, 5000 pneumonia images, and 4420 COVID-19 images was collected from eleven different publicly available datasets by Badwini et. al. [1] Using data augmentation with the Augmentor package, we generate an additional 580 COVID-19 images to construct a more balanced dataset.

dir = os.getcwd()
p = Augmentor.Pipeline(os.path.join(dir, 'preXrayData/covid'))
p.rotate(probability=0.7, max_left_rotation=10, max_right_rotation=10) # Rotate 70% of the images between 10 and -10 degrees
p.random_distortion(probability=1, grid_width=4, grid_height=4, magnitude=8) # Randomly distort the images while maintaining their aspect ratio
p.flip_left_right(probability=1) # Mirror the images from left to right
p.process()
p.sample(580)

Initialised with 4999 image(s) found.

Processing <PIL.Image.Image image mode=L size=3050x2539 at 0x7F103893BF40>: 100%|██████████| 4999/4999 [06:11<00:00, 13.46 Samples/s]   
Processing <PIL.Image.Image image mode=L size=299x299 at 0x7F0FFC56D810>: 100%|██████████| 580/580 [00:49<00:00, 11.61 Samples/s]     

'''Visualize augmented image.'''

aug_img_path = 'preXrayData/covid/covid_original_COVID-597.png_b8c0f5bb-77e0-4749-9359-d746af324bcb.png'
display(Image.open(aug_img_path))

'''Move additional generated images into correct folder.'''

src_path = os.path.join(dir, 'preXrayData/covid/output')
for each_file in Path(src_path).glob('*.*'): # grabs all files
    tar_path = each_file.parent.parent
    each_file.rename(tar_path.joinpath(each_file.name))
os.remove(src_path)

'''Split the training and validation data.'''

data = os.path.join(dir, 'preXrayData')
splitfolders.ratio(data, output='postXrayData', seed=1337, ratio=(0.8,0.2))

Copying files: 20578 files [00:10, 1920.93 files/s]

The Model

InceptionV3 is an image recognition model that has been pre-trained on the ImageNet dataset, obtaining an accuracy of greater than 78.1%. Based on the paper Rethinking the Inception Architecture for Computer Vision by Szegedy et. al [2], the aim of the model was to demonstrate that although increased model size and computational cost tend to translate to immediate quality gains for most tasks - as long as enough labeled data is provided for training - computational efficiency and low parameter count are also valuable enabling factors. Thus, InceptionV3 focused on scaling up previous Inception iterations as efficiently as possible by using suitably factorized convolutions and aggressive regularization.

Why InceptionV3?

Why Transfer Learning?

Transfer learning is a method to overcome insufficient data and/or training resources by adding a head or final layer to a pre-trained network (replacing the original classifier). This involves taking a pre-trained model, extracting one of the layers, and using it as the input layer to a series of dense layers. A dense layer is a simple layer of neurons in which each neuron recieves input from all the neurons of the previous layer, thus caputuring complex patterns.

Model Architecture