load_fetal_ecg

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Dec 31 15:33:42 2019

@author: Dani Kiyasseh
"""

import os
import numpy as np
import wfdb
import pickle

#basepath = '/home/scro3517/Desktop/non-invasive-fetal-ecg-arrhythmia-database-1.0.0'

files = os.listdir(basepath)
files = [file for file in files if not os.path.isdir(file)]
files = [file.split('.hea')[0] for file in files if '.hea' in file]

abdomen = 'Abdomen_1'
samples = 2500

inputs = dict()
outputs = dict()
all_labels = []
for file in files:
    data = wfdb.rdsamp(os.path.join(basepath,file))
    signals = data[0]
    info = data[1]
    signal_indices = np.where([abd in abdomen for abd in info['sig_name']])[0]
    signals = signals.transpose()
    signal = signals[signal_indices]
    label = file.split('_')[0]
    
    inputs[file] = []
    outputs[file] = []
    nframes = signal.shape[1]//samples
    for nframe in range(nframes):
        frame = signal[0,nframe*samples:(nframe+1)*samples]
        inputs[file].append(frame)
        outputs[file].append(label)
        all_labels.append(label)
        
    inputs[file] = np.array(inputs[file])
    outputs[file] = np.array(outputs[file])
    
#%%
""" Retrieve Unique Class Names """
unique_labels = []
for label in all_labels:
    if label not in unique_labels:
        unique_labels.append(label)

""" Convert Names to Labels """
from sklearn.preprocessing import LabelEncoder
label_encoder = LabelEncoder()
label_encoder.fit(unique_labels)
for patient_number,labels in outputs.items():
    outputs[patient_number] = label_encoder.transform(labels)

#%%
""" Make New Directory to Avoid Contamination """
savepath = os.path.join(basepath,'patient_data',abdomen)
try:
    os.chdir(savepath)
except:
    os.makedirs(savepath)
#%%
""" Save Inputs and Labels Dicts For Splitting Later """
with open(os.path.join(savepath,'ecg_signal_frames_fetal.pkl'),'wb') as f:
    pickle.dump(inputs,f)
with open(os.path.join(savepath,'ecg_signal_labels_fetal.pkl'),'wb') as f:
    pickle.dump(outputs,f)