Arrythmia prediction on ECG data using CNN¶
Credit: AITS Cainvas Community
Photo by Chan Luu on Behance, Adobe
The use of deep learning models in medical fields can help reduce error rates and increase the possibility of an earlier diagnosis for better treatement.
Dataset¶
Data source: Physionet's MIT-BIH Arrhythmia Dataset
The signals in the dataset correspond to electrocardiogram (ECG) shapes of heartbeats for the normal case and the cases affected by different arrhythmias and myocardial infarction. These signals are preprocessed and segmented, with each segment corresponding to a heartbeat.
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.metrics import f1_score, confusion_matrix
from sklearn.utils import resample
import tensorflow.keras
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv1D, MaxPooling1D, Dense, Dropout, Flatten, MaxPool1D, Convolution1D
from tensorflow.keras.layers import BatchNormalization
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
import random
train = pd.read_csv('https://cainvas-static.s3.amazonaws.com/media/user_data/cainvas-admin/mitbih_train.csv',header=None)
test = pd.read_csv('https://cainvas-static.s3.amazonaws.com/media/user_data/cainvas-admin/mitbih_test.csv',header=None)
train
train.describe()
The attribute values are almost in the same range.
The classes¶
There are five classes in the dataset:
- 0 - Non-ecotic beats (normal beat)
- 1 - Supraventricular ectopic beats
- 2 - Ventricular ectopic beats
- 3 - Fusion beats
- 4 - Unknown beats
# The classes
label_names = ['Non-ecotic beats (normal beat)', 'Supraventricular ectopic beats', 'Ventricular ectopic beats', 'Fusion beats', 'Unknown beats']
labels = train[187].astype('int64') # last column has the labels
print("Count in each label: ")
print(labels.value_counts())
plt.barh(list(set(labels)), list(labels.value_counts()))
The dataset is very imbalanced.
The samples have to be separated into categories before resampling to provide a balanced dataset
# Separating the train dataframe into 5 individual ones based on class labels, and sampling 50000 from each.
train_lbl0 = resample(train[train[187]==0], replace=True, n_samples=50000, random_state=113)
train_lbl1 = resample(train[train[187]==1], replace=True, n_samples=50000, random_state=113)
train_lbl2 = resample(train[train[187]==2], replace=True, n_samples=50000, random_state=113)
train_lbl3 = resample(train[train[187]==3], replace=True, n_samples=50000, random_state=113)
train_lbl4 = resample(train[train[187]==4], replace=True, n_samples=50000, random_state=113)
# Concatenate the 5 dataframes into 1
train = pd.concat([train_lbl0, train_lbl1, train_lbl2, train_lbl3, train_lbl4])
labels = train[187].astype('int64') # last column has the labels
print("Count in each label: ")
print(labels.value_counts())
Visualization¶
plt.plot(np.array(train_lbl0.sample(1))[0, :187])
plt.title(label_names[0])
plt.plot(np.array(train_lbl1.sample(1))[0, :187])
plt.title(label_names[1])
plt.plot(np.array(train_lbl2.sample(1))[0, :187])
plt.title(label_names[2])
plt.plot(np.array(train_lbl3.sample(1))[0, :187])
plt.title(label_names[3])
plt.plot(np.array(train_lbl4.sample(1))[0, :187])
plt.title(label_names[4])
Preprocessing¶
# Adding some noise to increase efficiency of the trained model
def gaussian_noise(signal):
noise = np.random.normal(0,0.05,187)
return signal + noise
# Visualization with added noise
sample = train_lbl0.sample(1).values[0]
sample_with_noise = gaussian_noise(sample[:187])
plt.subplot(1, 1, 1)
plt.plot(sample[:187])
plt.plot(sample_with_noise)
# One hot encoding the output of the model
ytrain = tensorflow.keras.utils.to_categorical(train[187])
ytest = tensorflow.keras.utils.to_categorical(test[187])
# Input to the model
xtrain = train.values[:, :187]
xtest = test.values[:, :187]
# Adding noise
for i in range(xtrain.shape[0]):
xtrain[i, :187] = gaussian_noise(xtrain[i, :187])
# Viewing the shapes
xtrain = np.expand_dims(xtrain, 2)
xtest = np.expand_dims(xtest, 2)
print("Shape of training data: ")
print("Input: ", xtrain.shape)
print("Output: ", ytrain.shape)
print("\nShape of test data: ")
print("Input: ", xtest.shape)
print("Output: ", ytest.shape)
The model¶
model = Sequential()
model.add(Conv1D(64, 6, activation = 'relu', input_shape = xtrain[0].shape))
model.add(MaxPool1D(3, 2))
model.add(Conv1D(64, 6, activation = 'relu'))
model.add(MaxPool1D(3, 2))
model.add(Conv1D(64, 6, activation = 'relu'))
model.add(MaxPool1D(3, 2))
model.add(Flatten())
model.add(Dense(64, activation = 'relu'))
model.add(Dense(32, activation = 'relu'))
model.add(Dense(5, activation = 'softmax'))
model.compile(optimizer = tensorflow.keras.optimizers.Adam(0.001), loss = 'categorical_crossentropy', metrics = ['accuracy'])
model.summary()
history = model.fit(xtrain, ytrain, epochs = 8, batch_size = 32, validation_data = (xtest, ytest))
Plotting the metrics¶
def plot(history, variable, variable2):
plt.plot(range(len(history[variable])), history[variable])
plt.plot(range(len(history[variable2])), history[variable2])
plt.legend([variable, variable2])
plt.title(variable)
plot(history.history, "accuracy", "val_accuracy")
plot(history.history, "loss", "val_loss")
model.save('ecg_arryhtmia.h5')
Model evaluation¶
ypred = model.predict(xtest)
cm = confusion_matrix(ytest.argmax(axis=1), ypred.argmax(axis=1))
cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
for i in range(cm.shape[1]):
for j in range(cm.shape[0]):
plt.text(j, i, format(cm[i, j], '.2f'), horizontalalignment="center", color="black")
plt.imshow(cm, cmap=plt.cm.Blues)
Seems like our model is performing really well.
# Test data class labels spread
print("The distribution of test set labels")
print(test[187].value_counts())
print('F1_score = ', f1_score(ytest.argmax(axis=1), ypred.argmax(axis=1), average = 'macro'))
Given that our test set is imbalanced, the high f1-score indicates that our model has good performnace.
Prediction¶
i = random.randint(0, len(xtest)-1)
output = model(np.expand_dims(xtest[i], 0))
pred = output.numpy()[0]
plt.plot(xtest[0])
print("Actual label: ", label_names[np.argmax(ytest[i])])
print("Model prediction : ", label_names[np.argmax(pred)], " with probability ", pred[np.argmax(pred)])
deepC¶
!deepCC ecg_arryhtmia.h5
i = random.randint(0, len(xtest)-1)
np.savetxt('sample.data', (xtest[i]).flatten())
!ecg_arryhtmia_deepC/ecg_arryhtmia.exe sample.data
pred = np.loadtxt('deepSea_result_1.out')
plt.plot(xtest[0])
print("Actual label: ", label_names[np.argmax(ytest[i])])
print("Model prediction : ", label_names[np.argmax(pred)], " with probability ", pred[np.argmax(pred)])