butter bandpass filter added

preprocess_and_segmentation.py

@@ -3,15 +3,25 @@ from scipy.io import loadmat
 import os
 import re
 import matplotlib.pyplot as plt
+from preprocessor import butter_filter
+
+data_dir = 'sample_of_data/Training_WFDB'
+
+fs = 500  # Hz
+lead_labels = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']
+
+# filter
+lf_filter = 0.5  # Hz
+hf_filter = 20  # Hz
+order_filter = 4
 
 
 def load_data(data_dir):
-    """ Load the subject signals (".mat" files) and info (".hea" files)
-        into a dictionary """
+    """ Load and preprocess the subject signals (".mat" files) and
+    info (".hea" files) into a dictionary """
 
     # signals (".mat" files)
     data = {}
-    lead_labels = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']
     for root, dirs, files in os.walk(data_dir):
         for file in files:
             if file.endswith(".mat"):
@@ -24,6 +34,10 @@ def load_data(data_dir):
 
                 # add the signals to the data dictionary
                 signals = loadmat(os.path.join(root, file))['val']  # load signals from the mat file
+
+                # filter signals (butter bandpass)
+                signals = butter_filter(signals, lf=lf_filter, hf=hf_filter, fs=fs, order=order_filter)
+
                 for i in range(signals.shape[0]):
                     data[subj_id][lead_labels[i]] = signals[i, :]
 
@@ -105,9 +119,8 @@ def segment_all_dict_data(data_dict, seg_width, overlap_perc):
     return segmented_dict_of_data
 
 
-data_dir = 'sample_of_data/Training_WFDB'
-
 data = load_data(data_dir)
+
 data_segmented = segment_all_dict_data(data, 500, 0.5)
 
 subj1_leadI = data_segmented['A0001']['I']

preprocessor.py

+import scipy.signal as signal
+
+FS = 500
+
+
+def butter_filter(data, lf=1, hf=40, fs=FS, order=2):
+    wbut = [ 2 * lf/fs, 2 * hf / fs]
+    bbut, abut = signal.butter(order, wbut, btype='bandpass')
+
+    if type(data) == dict:
+        for key in data:
+            data[key] = signal.filtfilt(bbut, abut, data[key])
+        return data
+    else:
+        return signal.filtfilt(bbut, abut, data)
\ No newline at end of file