Add: included butter filter

app.py

@@ -3,10 +3,6 @@ from make_predictions import make_prediction
 import gradio as gr
 
 def predict(file):
-    # print(file.name)
-    # print(file)
-    with open(file.name, 'rb') as f:
-        print(f.read())
     print("=====================================")
     ans = make_prediction(file.name)
     print(ans)

baseline_wander_removal.py

@@ -1,6 +1,6 @@
 import numpy as np
 from scipy.signal import medfilt
-# import matplotlib.pyplot as plt
+import matplotlib.pyplot as plt
 # import seaborn as sns
 
 # sns.set_theme()
@@ -80,3 +80,19 @@ def band_pass_filter(signal):
     # result = result/max_val
 
     return result
+
+
+if __name__ == '__main__':
+    signal = np.loadtxt("./Lead 2/n1_lead2.txt")[:1000]
+    signal_bs = bw_remover(BASIC_SRATE, signal)
+    signal_flt = band_pass_filter(signal_bs)
+    plt.figure(figsize=(16, 9))
+    plt.subplot(2, 1, 1)
+    plt.plot(signal)
+    plt.grid(True)
+    plt.title("RAW signal")
+    plt.subplot(2, 1, 2)
+    plt.plot(signal_flt)
+    plt.title("baseline removed signal")
+    plt.grid(True)
+    plt.show()

make_predictions.py

@@ -1,22 +1,17 @@
 import numpy as np
-from read_data import prepare_all_leads
-from postprocessing import make_predictions_indi, labels_map
+from read_data import prepare_all_leads, visualize_sig
+from postprocessing import predict_disease, labels_map
 from scipy import stats as st
 import warnings
 
 warnings.filterwarnings("ignore")
 
-def make_prediction(path):
-    leads = prepare_all_leads(path)
-    # print(leads[0][0].shape)
-    # print(leads[1][0].shape)
-    # print(leads[2][0].shape)
-    # visualize_sig([leads[0][0], leads[1][0], leads[2][0]])
-    x = make_predictions_indi(*leads)
-    # print(x.mean(axis=0))
-    # print(np.argmax(x, axis=1))
-    index = st.mode(np.argmax(x, axis=1))[0][0]
-    confidence = x.mean(axis=0)[index]
+def make_prediction(path, visualize=False, butter_filter=True):
+    leads = prepare_all_leads(path, butter_filter=butter_filter)
+    visualize_sig([leads[0][0], leads[1][0], leads[2][0]]) if visualize else None
+    all, x, y, z = predict_disease(*leads)
+    index = st.mode(np.argmax(all, axis=1))[0][0]
+    confidence = all.mean(axis=0)[index]
     return labels_map[index], float(confidence)
 
 if __name__ == "__main__":

model/{ResNet-lead-0.pth → seven-diseases/ResNet-lead-0-BEST.pth}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f573373d1fb5d86be738594aebd2ac0a42dde9c3358da27324277cc88ce0fd03
-size 2081879
+oid sha256:33a3ffa62adedac7bfeb52161baaff9ebd3532663ddb6b25818e97a76def0c75
+size 2056539

model/{ResNet-lead-1.pth → seven-diseases/ResNet-lead-1-BEST.pth}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e649f7c20301dfd614f1504ff8c23da284a1e422acbdbb74f63acbece25c3f97
-size 2081879
+oid sha256:d7180c82754cd4f66f378204a29e9bd8b0b2c628495673114e68a55e0e31fac8
+size 2056539

model/{ResNet-lead-2.pth → seven-diseases/ResNet-lead-2-BEST.pth}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bb6a7635539ee96d4dd6011afb46e06a94d54a49e2f70234cfbfb990f1149918
-size 2081879
+oid sha256:167cd4161b3fac862bedfce0132e13cd4c3947834b89d26bfe55a49db319fe91
+size 2056539

model_def.py

@@ -1,28 +1,37 @@
 import torch
 from torch import nn
 import torch.nn.functional as F
-
-
 class NeuralNetwork(nn.Module):
     def __init__(self):
         super().__init__()
         n_filters = 64
         self.conv_1 = nn.Conv1d(        1, n_filters, 8, stride=1, padding='same')
+        self.norm_1 = nn.BatchNorm1d(n_filters)
         self.conv_2 = nn.Conv1d(n_filters, n_filters, 5, stride=1, padding='same')
+        self.norm_2 = nn.BatchNorm1d(n_filters)
         self.conv_3 = nn.Conv1d(n_filters, n_filters, 3, stride=1, padding='same')
+        self.norm_3 = nn.BatchNorm1d(n_filters)
         self.conv_4 = nn.Conv1d(        1, n_filters, 1, stride=1, padding='same') # Expanding for addition
+        self.norm_4 = nn.BatchNorm1d(n_filters)
 
         self.conv_5 = nn.Conv1d(  n_filters, n_filters*2, 8, stride=1, padding='same')
+        self.norm_5 = nn.BatchNorm1d(n_filters*2)
         self.conv_6 = nn.Conv1d(n_filters*2, n_filters*2, 5, stride=1, padding='same')
+        self.norm_6 = nn.BatchNorm1d(n_filters*2)
         self.conv_7 = nn.Conv1d(n_filters*2, n_filters*2, 3, stride=1, padding='same')
+        self.norm_7 = nn.BatchNorm1d(n_filters*2)
         self.conv_8 = nn.Conv1d(  n_filters, n_filters*2, 1, stride=1, padding='same')
-        
+        self.norm_8 = nn.BatchNorm1d(n_filters*2)
+
         self.conv_9  = nn.Conv1d(n_filters*2, n_filters*2, 8, stride=1, padding='same')
+        self.norm_9 = nn.BatchNorm1d(n_filters*2)
         self.conv_10 = nn.Conv1d(n_filters*2, n_filters*2, 5, stride=1, padding='same')
+        self.norm_10 = nn.BatchNorm1d(n_filters*2)
         self.conv_11 = nn.Conv1d(n_filters*2, n_filters*2, 3, stride=1, padding='same')
-        self.conv_12 = nn.Conv1d(n_filters*2, n_filters*2, 1, stride=1, padding='same')
-
-        self.classifier = nn.Linear(128, 5)
+        self.norm_11 = nn.BatchNorm1d(n_filters*2)
+#         self.conv_12 = nn.Conv1d(n_filters*2, n_filters*2, 1, stride=1, padding='same')
+        self.norm_12 = nn.BatchNorm1d(n_filters*2)
+        self.classifier = nn.Linear(128, 7)
         self.log_softmax = nn.LogSoftmax(dim=1)
 
     def forward(self, x):
@@ -30,38 +39,57 @@ class NeuralNetwork(nn.Module):
         
         # Block 1
         a = self.conv_1(x)
+        a = self.norm_1(a)
         a = F.relu(a)
+        
         b = self.conv_2(a)
+        b = self.norm_2(b)
         b = F.relu(b)
+        
         c = self.conv_3(b)
+        c = self.norm_3(c)
+        
         shortcut = self.conv_4(x)
+        shortcut = self.norm_4(shortcut)
+        
         output_1 = torch.add(c, shortcut)
         output_1 = F.relu(output_1)
         
         #Block 2
         a = self.conv_5(output_1)
+        a = self.norm_5(a)
         a = F.relu(a)
+        
         b = self.conv_6(a)
+        b = self.norm_6(b)
         b = F.relu(b)
+        
         c = self.conv_7(b)
+        c = self.norm_7(c)
         shortcut = self.conv_8(output_1)
+        shortcut = self.norm_8(shortcut)
+        
         output_2 = torch.add(c, shortcut)
         output_2 = F.relu(output_2)
         
         #Block 3
         a = self.conv_9(output_2)
+        a = self.norm_9(a)
         a = F.relu(a)
+        
         b = self.conv_10(a)
+        b = self.norm_10(b)
         b = F.relu(b)
+        
         c = self.conv_11(b)
-        shortcut = self.conv_12(output_2)
+        c = self.norm_11(c)
+        
+#         shortcut = self.conv_12(output_2)
+        shortcut = self.norm_12(shortcut)
+        
         output_3 = torch.add(c, shortcut)
         output_3 = F.relu(output_3)
-        res = self.classifier(output_3.mean((2,)))
-        logits = self.log_softmax(res)
-        return logits
-
-
-if __name__ == '__main__':
-    model = NeuralNetwork()
-    print(model)
+        
+        logits = self.classifier(output_3.mean((2,)))
+        res = self.log_softmax(logits)
+        return res

postprocessing.py

@@ -4,23 +4,33 @@ from torch import nn
 from model_def import NeuralNetwork
 
 labels_map = {
-    0 : "atrial fibrillation",
-    1 : "sinus arrhythmia",
-    2 : "bradycardia",
-    3 : "1st degree av block",
-    4 : "sinus rhythm",
+    0: 'sinus_rhythm',
+    1: 'atrial_fibrillation',
+    2: 'av_block',
+    3: 'bradycardia',
+    4: 'sinus_arrhythmia',
+    5: 'sinus_rhythm-sinus_arrhythmia',
+    6: 'sinus_rhythm-av_block'
 }
 
-PATH = 'model/' #ResNet-lead-0.pth'
+PATH = 'model/seven-diseases/' #ResNet-lead-0.pth'
+lead_1 = f"{PATH}/ResNet-lead-0-BEST.pth"
+lead_2 = f"{PATH}/ResNet-lead-1-BEST.pth"
+lead_3 = f"{PATH}/ResNet-lead-2-BEST.pth"
+
+# PATH = 'model/' #ResNet-lead-0.pth'
+# lead_1 = f"{PATH}/ResNet-lead-0.pth"
+# lead_2 = f"{PATH}/ResNet-lead-1.pth"
+# lead_3 = f"{PATH}/ResNet-lead-2.pth"
 
 lead_1_model = NeuralNetwork()
-lead_1_model.load_state_dict(torch.load(f"{PATH}/ResNet-lead-0.pth", map_location=torch.device('cpu')))
+lead_1_model.load_state_dict(torch.load(lead_1, map_location=torch.device('cpu')))
 
 lead_2_model = NeuralNetwork()
-lead_2_model.load_state_dict(torch.load(f"{PATH}/ResNet-lead-1.pth", map_location=torch.device('cpu')))
+lead_2_model.load_state_dict(torch.load(lead_2, map_location=torch.device('cpu')))
 
 lead_3_model = NeuralNetwork()
-lead_3_model.load_state_dict(torch.load(f"{PATH}/ResNet-lead-2.pth", map_location=torch.device('cpu')))
+lead_3_model.load_state_dict(torch.load(lead_3, map_location=torch.device('cpu')))
 
 def helper(sig, model):
     inpt = sig[:, np.newaxis, :]
@@ -30,9 +40,14 @@ def helper(sig, model):
     prediction_scores = res.numpy()
     return prediction_scores
 
-def make_predictions_indi(lead1, lead2, lead3):
+def predict_disease(lead1, lead2, lead3):
     p1 = helper(lead1, lead_1_model)
     p2 = helper(lead2, lead_2_model)
     p3 = helper(lead3, lead_3_model)
+    # print(p1.argmax(axis=1))
+    # print(p2.argmax(axis=1))
+    # print(p3.argmax(axis=1))
     p_avg = (p1 + p2 + p3)/3
-    return p_avg
+    # print(p_avg.argmax(axis=1))
+    # print("-----------------")
+    return p_avg, p1, p2, p3

read_data.py

@@ -1,53 +1,86 @@
 import numpy as np
-from scipy.signal import resample
+from scipy.signal import resample, butter, filtfilt
 from baseline_wander_removal import bw_remover
 import pywt
-# import matplotlib.pyplot as plt
-import codecs
+import matplotlib.pyplot as plt
 
-def normalize(sig):
-    return 2*((sig-np.min(sig))/(np.max(sig)-np.min(sig)))
+def normalize(sig, val=2):
+    return val*((sig-np.min(sig))/(np.max(sig)-np.min(sig)))
 
-def prepare_all_leads(path):
+def butter_lowpass_filter(data, cutoff, fs, order):
+    nyq = 0.5 * fs
+    normal_cutoff = cutoff / nyq
+    # Get the filter coefficients 
+    b, a = butter(order, normal_cutoff, btype='low', analog=False)
+    y = filtfilt(b, a, data)
+    return y
+
+def visualize_sig(sig, filename="test.png"):
+    fig, ax = plt.subplots(3, 1, figsize=(10, 10))
+    ax[0].plot(sig[0])
+    ax[1].plot(sig[1])
+    ax[2].plot(sig[2])
+    # plt.show()
+    plt.savefig(filename)
+
+def preprocess_one_chunk(chunk, lvl=2):
+    chunk = resample(chunk, 1000)
+    chunk = pywt.wavedec(chunk, 'db6', level=lvl)[0]
+    # print(x.shape)
+    # x = pad(x)
+    chunk = normalize(chunk)
+    return chunk
+
+def prepare_all_leads(path, butter_filter=False):
     if path.endswith(".txt"):
-        sig = np.loadtxt(path, delimiter=',', unpack=True)
+        signal = np.loadtxt(path, delimiter=',', unpack=True)
     elif path.endswith(".npy"):
         sig = np.load(path, allow_pickle=True)
         x = pywt.wavedec(sig[0], 'db6', level=2)[0]
         y = pywt.wavedec(sig[1], 'db6', level=2)[0]
         z = pywt.wavedec(sig[2], 'db6', level=2)[0]
-        # visualize_sig([x, y, z])
         return x[None, :], y[None, :], z[None, :]
-    # return sig
-    freq = sig.shape[1] // 60
+    freq = signal.shape[1] // 60
+    if freq < 250:
+        lvl = 1
+    else:
+        lvl = 2
+    sig = [signal[0], signal[1], signal[2]]
+
     sig[0] = bw_remover(freq, sig[0])
     sig[1] = bw_remover(freq, sig[1])
     sig[2] = bw_remover(freq, sig[2])
 
-    # visualize_sig(sig)
+    if butter_filter:
+        cutoff = 20      # desired cutoff frequency of the filter, Hz ,      slightly higher than actual 1.2 Hz
+        order = 1       # sin wave can be approx represented as quadratic
+
+        sig[0] = butter_lowpass_filter(sig[0], cutoff, freq, order)
+        sig[1] = butter_lowpass_filter(sig[1], cutoff, freq, order)
+        sig[2] = butter_lowpass_filter(sig[2], cutoff, freq, order)
+
     sig_length = freq*2
-    total_samples = sig[0].shape[0] // sig_length
+    total_samples = sig[0].shape[0] // 1000
 
     lead_1 = []
     lead_2 = []
     lead_3 = []
     for i in range(total_samples):
         x = sig[0][i*sig_length:(i+1)*sig_length]
-        x = pywt.wavedec(x, 'db6', level=1)[0]
-        x = resample(x, 259)
-        x = normalize(x)
-        lead_1.append(x)
-
         y = sig[1][i*sig_length:(i+1)*sig_length]
-        y = pywt.wavedec(y, 'db6', level=1)[0]
-        y = resample(y, 259)
-        y = normalize(y)
-        lead_2.append(y)
-
         z = sig[2][i*sig_length:(i+1)*sig_length]
-        z = pywt.wavedec(z, 'db6', level=1)[0]
-        z = resample(z, 259)
-        z = normalize(z)
+        
+        x = preprocess_one_chunk(x, lvl=lvl)
+        y = preprocess_one_chunk(y, lvl=lvl)
+        z = preprocess_one_chunk(z, lvl=lvl)
+        
+        lead_1.append(x)
+        lead_2.append(y)
         lead_3.append(z)
-    
+
     return np.asarray(lead_1), np.asarray(lead_2), np.asarray(lead_3)
+
+
+def pad(sig):
+    sig = np.pad(sig, (0, 258-sig.shape[0]), 'constant')
+    return sig