AI-Powered Pneumonia Detection Using Chest X-rays
In the age of artificial intelligence, machine learning models have become powerful tools in the field of medical diagnostics. One of the areas where AI has shown great promise is in the detection of pneumonia from chest X-rays. In this blog post, I will walk you through my project where I developed different convolutional neural networks (CNNs) to accurately distinguish between healthy lungs and those affected by pneumonia.
Project Overview
The goal of this project was to leverage the TensorFlow library to create various CNN models for pneumonia detection. I used a combination of a custom simple CNN and more complex models through transfer learning with pre-trained MobileNet and VGG16 architectures.
Dataset
The dataset used in this project consists of chest X-ray images categorized into two classes: Normal and Pneumonia. The dataset was split into training, validation, and test sets to ensure robust evaluation of the models. Download the dataset from kaggle using this link: https://www.kaggle.com/datasets/paultimothymooney/chest-xray-pneumonia
Models and Methodology
1. Custom Simple CNN
I started with a custom simple CNN model to establish a baseline performance. The architecture consists of several convolutional layers followed by a dense layer with a softmax activation to output the class probabilities.
import tensorflow as tf
from tensorflow.keras.layers import Activation, Dense, Flatten, Conv2D
from tensorflow.keras.models import Sequential
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.optimizers import Adam
train_path = 'balanced_chest_xray_dataset/train'
valid_path = 'balanced_chest_xray_dataset/valid'
test_path = 'balanced_chest_xray_dataset/test'
data_generator = ImageDataGenerator(preprocessing_function=tf.keras.applications.mobilenet.preprocess_input)
train_batches = data_generator.flow_from_directory(directory=train_path, target_size=(224,224), batch_size=10)
valid_batches = data_generator.flow_from_directory(directory=valid_path, target_size=(224,224), batch_size=10)
test_batches = data_generator.flow_from_directory(directory=test_path, target_size=(224,224), batch_size=10, shuffle=False)
simple_cnn = Sequential([
Conv2D(16, kernel_size=(3,3), input_shape=(224,224,3), activation='relu', padding='same'),
Conv2D(32, kernel_size=(3,3), activation='relu', padding='same'),
Conv2D(64, kernel_size=(3,3), activation='relu', padding='same'),
Conv2D(128, kernel_size=(3,3), activation='relu', padding='same'),
Flatten(),
Dense(2, activation='softmax')
])
simple_cnn.compile(optimizer=Adam(learning_rate=0.0001), loss='categorical_crossentropy', metrics=['accuracy'])
simple_cnn.fit(x=train_batches, steps_per_epoch=len(train_batches), validation_data=valid_batches, validation_steps=len(valid_batches), epochs=25, verbose=2)
simple_cnn.save('models/simple_cnn_balanced_dataset.h5')2. Transfer Learning with MobileNet
MobileNet, known for its efficiency and small size, was used to create multiple models by fine-tuning different layers. This approach allows us to leverage the pre-trained weights on ImageNet and adapt them to our specific task.
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import Model
train_path = '/content/drive/MyDrive/Machine_Learning_Projects/equal_chest_xray_dataset/train'
valid_path = '/content/drive/MyDrive/Machine_Learning_Projects/equal_chest_xray_dataset/valid'
test_path = '/content/drive/MyDrive/Machine_Learning_Projects/equal_chest_xray_dataset/test'
data_generator = ImageDataGenerator(preprocessing_function=tf.keras.applications.mobilenet.preprocess_input)
train_batches = data_generator.flow_from_directory(directory=train_path, target_size=(224,224), batch_size=10)
valid_batches = data_generator.flow_from_directory(directory=valid_path, target_size=(224,224), batch_size=10)
test_batches = data_generator.flow_from_directory(directory=test_path, target_size=(224,224), batch_size=10, shuffle=False)
train_layers = [5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86]
for n in train_layers:
mobile = tf.keras.applications.mobilenet.MobileNet()
x = mobile.layers[-6].output
output = Dense(units=2, activation='softmax')(x)
model = Model(inputs=mobile.input, outputs=output)
for layer in model.layers[:-n]:
layer.trainable = False
model.compile(optimizer=Adam(learning_rate=0.0001), loss='categorical_crossentropy', metrics=['accuracy'])
model.fit(x=train_batches, steps_per_epoch=len(train_batches), validation_data=valid_batches, validation_steps=len(valid_batches), epochs=25, verbose=2)
model.save(f'/content/drive/MyDrive/Machine_Learning_Projects/models/mobilenet-full-{n}.h5')3. Transfer Learning with VGG16
VGG16 is a well-known deep learning model with a large size and high performance. By fine-tuning the VGG16 model, we can achieve higher accuracy in detecting pneumonia.
import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import Sequential
vgg16_model = tf.keras.applications.vgg16.VGG16()
vgg16_model.summary()
train_path = 'balanced_chest_xray_dataset/train'
valid_path = 'balanced_chest_xray_dataset/valid'
test_path = 'balanced_chest_xray_dataset/test'
data_generator = ImageDataGenerator(preprocessing_function=tf.keras.applications.mobilenet.preprocess_input)
train_batches = data_generator.flow_from_directory(directory=train_path, target_size=(224,224), batch_size=10)
valid_batches = data_generator.flow_from_directory(directory=valid_path, target_size=(224,224), batch_size=10)
test_batches = data_generator.flow_from_directory(directory=test_path, target_size=(224,224), batch_size=10, shuffle=False)
model = Sequential()
for layer in vgg16_model.layers[:-1]:
model.add(layer)
for layer in model.layers[:-20]:
layer.trainable = False
model.add(Dense(units=2, activation='softmax'))
model.compile(optimizer=Adam(learning_rate=0.0001), loss='categorical_crossentropy', metrics=['accuracy'])
model.fit(x=train_batches, steps_per_epoch=len(train_batches), validation_data=valid_batches, validation_steps=len(valid_batches), epochs=25, verbose=2)
model.save('models/balanced_chest_xray_vgg16.h5')Evaluation
After training, the models were evaluated on the test set to measure their performance. Confusion matrices and classification reports were generated to provide insights into the models’ accuracy and precision.
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from sklearn.metrics import confusion_matrix, classification_report
import itertools
import matplotlib.pyplot as plt
import numpy as np
import os
test_path = '/content/drive/MyDrive/Machine_Learning_Projects/equal_chest_xray_dataset/test'
data_generator = ImageDataGenerator(preprocessing_function=tf.keras.applications.mobilenet.preprocess_input)
test_batches = data_generator.flow_from_directory(directory=test_path, target_size=(224,224), batch_size=10, shuffle=False)
model_folder = os.listdir('/content/drive/MyDrive/Machine_Learning_Projects/models')
for model in model_folder:
fine_tuned_model = load_model(f'/content/drive/MyDrive/Machine_Learning_Projects/models/{model}')
test_imgs, test_labels = next(test_batches)
predictions = fine_tuned_model.predict(x=test_batches, steps=len(test_batches), verbose=0)
cm = confusion_matrix(y_true=test_batches.classes, y_pred=np.argmax(predictions, axis=-1))
def plot_confusion_matrix(cm, classes, normalize=False, title=f'{model}', cmap=plt.cm.Blues):
plt.imshow(cm, interpolation='nearest', cmap=cmap)
plt.title(title)
plt.colorbar()
tick_marks = np.arange(len(classes))
plt.xticks(tick_marks, classes, rotation=45)
plt.yticks(tick_marks, classes)
if normalize:
cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
print('Normalized confusion matrix')
else:
print(f'{model}')
print(cm)
thresh = cm.max() / 2
for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
plt.text(j, i, cm[i, j], horizontalalignment='center', color='white' if cm[i, j] > thresh else 'black')
plt.tight_layout()
plt.ylabel('True label')
plt.xlabel('Predicted label')
cm_plot_labels = ['Normal', 'Pneumonia']
plot_confusion_matrix(cm=cm, classes=cm_plot_labels, title=f'{model}')
report = classification_report(y_true=test_batches.classes, y_pred=np.argmax(predictions, axis=-1), target_names=['Normal', 'Pneumonia'])
print(report)Conclusion
This project demonstrates the power and potential of using AI and deep learning models for medical diagnostics. By leveraging custom CNNs and transfer learning with MobileNet and VGG16, we were able to create models capable of accurately distinguishing between healthy lungs and those affected by pneumonia using chest X-rays.
I hope you found this blog informative and inspiring. Your feedback is incredibly valuable to me, so please share your thoughts and suggestions in the comments below. If you’re interested in exploring the code and experimenting with the models yourself, you can visit my GitHub repository. Let’s collaborate and innovate together in the exciting field of AI-powered healthcare. Visit my github account for the full project: https://github.com/MargaretKagwiria/AI-Pneumonia-Detection-Using-Chest-X-rays/tree/main
