Covid-19 Diagnosis using Radiography Images

Classifying radiography images using a pre-trained ResNet-18 architecture

Last updated on Oct 31, 2021 3 min read Deep Learning, Image Classification

In this notebook, we will try to classify images from the Covid-19 Radiography Dataset[1] using a pre-trained ResNet-18 network. We will be using the PyTorch library for building our network.

Disclaimer: This model should be used only for learning purposes as Covid-19 diagnosis is an ongoing research topic.

Download the dataset

Firstly, you need to download the dataset from Kaggle. Check these steps for detailed instructions.

Data preprocessing

As a first, step download the dataset from Kaggle and create a new PyTorch dataset using the ImageFolder method.

Also, we are defining a transformer to Resize the images to 224x224 px and then converting the image to a Tensor.

	import matplotlib.pyplot as plt
	from torchvision import datasets
	import torchvision.transforms as transforms

	transform = transforms.Compose([transforms.Resize(size=(224,224)),
	transforms.ToTensor(),
	transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

	covid_19_dataset = datasets.ImageFolder('COVID-19 Radiography Database/',transform=transform)

view raw c0e7c573-aae9-4dd0-b58f-90f3f974d337 hosted with ❤ by GitHub

Next, we. split the dataset into training, validation and testing sets. We would be using 20% of the data for testing and 10% of the data for validation using a SubsetRandomSampler.

	from torch.utils.data import Dataset, DataLoader
	from torch.utils.data.sampler import SubsetRandomSampler
	import numpy as np


	batch_size = 4
	# percentage of training set to use as validation
	test_size = 0.2
	valid_size = 0.1

	#For test
	num_data = len(covid_19_dataset)
	indices_data = list(range(num_data))
	np.random.shuffle(indices_data)
	split_tt = int(np.floor(test_size * num_data))
	train_idx, test_idx = indices_data[split_tt:], indices_data[:split_tt]

	#For Valid
	num_train = len(train_idx)
	indices_train = list(range(num_train))
	np.random.shuffle(indices_train)
	split_tv = int(np.floor(valid_size * num_train))
	train_new_idx, valid_idx = indices_train[split_tv:],indices_train[:split_tv]


	# define samplers for obtaining training and validation batches
	train_sampler = SubsetRandomSampler(train_new_idx)
	test_sampler = SubsetRandomSampler(test_idx)
	valid_sampler = SubsetRandomSampler(valid_idx)

view raw 3c0a019c-6764-440e-a953-cde5fc33eb14 hosted with ❤ by GitHub

Next, lets define the DataLoader for training, testing and validation sets.

	import torch


	train_loader = torch.utils.data.DataLoader(covid_19_dataset, batch_size=batch_size,
	sampler=train_sampler, num_workers=1)
	valid_loader = torch.utils.data.DataLoader(covid_19_dataset, batch_size=batch_size,
	sampler=valid_sampler, num_workers=1)
	test_loader = torch.utils.data.DataLoader(covid_19_dataset, sampler = test_sampler, batch_size=batch_size,
	num_workers=1)
	classes = ['COVID-19','NORMAL','Viral Pneumonia']

view raw a761c07a-3e39-45c9-aa83-98cca8652378 hosted with ❤ by GitHub

Now, that our DataLoader is defined, lets use the train_loader to visualize a few images along with their classes.

	import matplotlib.pyplot as plt
	%matplotlib inline

	# helper function to un-normalize and display an image
	def imshow(img):
	img = img / 2 + 0.5 # unnormalize
	plt.imshow(np.transpose(img, (1, 2, 0))) # convert from Tensor image

	# obtain one batch of training images
	dataiter = iter(train_loader)
	images, labels = dataiter.next()
	images = images.numpy() # convert images to numpy for display

	# plot the images in the batch, along with the corresponding labels
	fig = plt.figure(figsize=(10, 4))
	# display 20 images
	for idx in np.arange(4):
	ax = fig.add_subplot(2, 10/2, idx+1, xticks=[], yticks=[])
	imshow(images[idx])
	ax.set_title(classes[labels[idx]])

view raw 962944c2-7c74-437d-a164-864525561eab hosted with ❤ by GitHub

png

Defining the model

Next, lets define the model. We will be using a pre-trained Resnet18 architecture for our classification task. As the number of images in our dataset is relatively less, using a transfer learning will be helpful.

	import torchvision.models as models
	model = models.resnet18(pretrained = True)

	train_on_gpu = torch.cuda.is_available()
	print(train_on_gpu)

view raw 3f257aea-c425-47f6-bff2-558c2cb511dc hosted with ❤ by GitHub

Downloading: "https://download.pytorch.org/models/resnet18-5c106cde.pth" to /root/.cache/torch/checkpoints/resnet18-5c106cde.pth



HBox(children=(FloatProgress(value=0.0, max=46827520.0), HTML(value='')))



True

	for param in model.parameters():
	param.requires_grad = False

view raw 2f87c54f-8f61-45c8-957c-cbea9cc10b1d hosted with ❤ by GitHub

Let’s add fully connected layers at the end of the network, to adjust the final layer’s output to correspond to the number of classes in our dataset.

	from torch import nn
	from collections import OrderedDict

	fc = nn.Sequential(OrderedDict([
	('fc1', nn.Linear(512,100)),
	('relu', nn.ReLU()),
	('fc2', nn.Linear(100,3)), # 3 is the number of classes we have in the dataset
	('output', nn.LogSoftmax(dim=1))
	]))

	model.fc = fc

view raw 28e1b9aa-4b01-444c-880e-e0bd73d33bdf hosted with ❤ by GitHub

	model

	if train_on_gpu:
	model.cuda()

view raw 30756667-f68c-4140-814d-27853e2a1056 hosted with ❤ by GitHub

We would be using CrossEntropyLoss as the loss function and a learning rate of 3e-5 for training.

	import torch.optim as optim

	# specify loss function
	criterion = nn.CrossEntropyLoss()

	# specify optimizer
	optimizer = optim.Adam(model.fc.parameters(), lr=3e-5)

view raw 497e5460-231e-40a9-b0ab-9c4ca8ad2e96 hosted with ❤ by GitHub

Training

Finally, lets train the network for 20 epochs

	# number of epochs to train the model
	n_epochs = 20 # you may increase this number to train a final model

	valid_loss_min = np.Inf # track change in validation loss

	for epoch in range(1, n_epochs+1):

	# keep track of training and validation loss
	train_loss = 0.0
	valid_loss = 0.0

	###################
	# train the model #
	###################
	model.train()
	for data, target in train_loader:
	# move tensors to GPU if CUDA is available
	if train_on_gpu:
	data, target = data.cuda(), target.cuda()
	# clear the gradients of all optimized variables
	optimizer.zero_grad()
	# forward pass: compute predicted outputs by passing inputs to the model
	output = model(data)
	# calculate the batch loss
	loss = criterion(output, target)
	# backward pass: compute gradient of the loss with respect to model parameters
	loss.backward()
	# perform a single optimization step (parameter update)
	optimizer.step()
	# update training loss
	train_loss += loss.item()*data.size(0)

	######################
	# validate the model #
	######################
	model.eval()
	for data, target in valid_loader:
	# move tensors to GPU if CUDA is available
	if train_on_gpu:
	data, target = data.cuda(), target.cuda()
	# forward pass: compute predicted outputs by passing inputs to the model
	output = model(data)
	# calculate the batch loss
	loss = criterion(output, target)
	# update average validation loss
	valid_loss += loss.item()*data.size(0)

	# calculate average losses
	train_loss = train_loss/len(train_loader.dataset)
	valid_loss = valid_loss/len(valid_loader.dataset)

	# print training/validation statistics
	print('Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f}'.format(
	epoch, train_loss, valid_loss))

	# save model if validation loss has decreased
	if valid_loss <= valid_loss_min:
	print('Validation loss decreased ({:.6f} --> {:.6f}). Saving model ...'.format(
	valid_loss_min,
	valid_loss))
	torch.save(model.state_dict(), 'model_nbl.pt')
	valid_loss_min = valid_loss

view raw fc555aff-cf98-43bf-9d8f-d639337b8925 hosted with ❤ by GitHub

Epoch: 1 	Training Loss: 0.569099 	Validation Loss: 0.045166
Validation loss decreased (inf --> 0.045166).  Saving model ...
Epoch: 2 	Training Loss: 0.447157 	Validation Loss: 0.035986
.
.
.
.
Epoch: 19 	Training Loss: 0.229328 	Validation Loss: 0.016049
Epoch: 20 	Training Loss: 0.248973 	Validation Loss: 0.013585
Validation loss decreased (0.014983 --> 0.013585).  Saving model ...

Evaluation

Now, that the training is complete, lets evaluate the performance of the network by predicting the classes for test dataset.

	# track test loss
	test_loss = 0.0
	class_correct = list(0. for i in range(3))
	class_total = list(0. for i in range(3))

	model.eval()
	i=1
	# iterate over test data
	len(test_loader)
	for data, target in test_loader:
	i=i+1
	if len(target)!=batch_size:
	continue

	# move tensors to GPU if CUDA is available
	if train_on_gpu:
	data, target = data.cuda(), target.cuda()
	# forward pass: compute predicted outputs by passing inputs to the model
	output = model(data)
	# calculate the batch loss
	loss = criterion(output, target)
	# update test loss
	test_loss += loss.item()*data.size(0)
	# convert output probabilities to predicted class
	_, pred = torch.max(output, 1)
	# compare predictions to true label
	correct_tensor = pred.eq(target.data.view_as(pred))
	correct = np.squeeze(correct_tensor.numpy()) if not train_on_gpu else np.squeeze(correct_tensor.cpu().numpy())
	# calculate test accuracy for each object class
	# print(target)

	for i in range(batch_size):
	label = target.data[i]
	class_correct[label] += correct[i].item()
	class_total[label] += 1

	# average test loss
	test_loss = test_loss/len(test_loader.dataset)
	print('Test Loss: {:.6f}\n'.format(test_loss))

	for i in range(3):
	if class_total[i] > 0:
	print('Test Accuracy of %5s: %2d%% (%2d/%2d)' % (
	classes[i], 100 * class_correct[i] / class_total[i],
	np.sum(class_correct[i]), np.sum(class_total[i])))
	else:
	print('Test Accuracy of %5s: N/A (no training examples)' % (classes[i]))

	print('\nTest Accuracy (Overall): %2d%% (%2d/%2d)' % (
	100. * np.sum(class_correct) / np.sum(class_total),
	np.sum(class_correct), np.sum(class_total)))

view raw 45db135d-f3e8-46db-8cde-b4b2763aac58 hosted with ❤ by GitHub

Test Loss: 0.037036

Test Accuracy of COVID-19: 100% (48/48)
Test Accuracy of NORMAL: 97% (268/276)
Test Accuracy of Viral Pneumonia: 91% (233/256)

Test Accuracy (Overall): 94% (549/580)

Conclusion

We used a pre-trained ResNet18 architecture for classifying images from the Radiography dataset. Our model gives a good accuracy of around 94% without reinventing the wheel. This demonstrates the power of transfer learning and shows how new problems can be tackled using existing models pre-trained on much larger datasets. With some minor enhancements like data augmentation, the accuracy can further be improved.

References

[1] M.E.H. Chowdhury, T. Rahman, A. Khandakar, R. Mazhar, M.A. Kadir, Z.B. Mahbub, K.R. Islam, M.S. Khan, A. Iqbal, N. Al-Emadi, M.B.I. Reaz, “Can AI help in screening Viral and COVID-19 pneumonia?” arXiv preprint, 29 March 2020, https://arxiv.org/abs/2003.13145. https://www.kaggle.com/tawsifurrahman/covid19-radiography-database

Deep Learning Image Classification