Datasets
Models
Downloads: 1
Diff of
/Inference/utils.py
[000000]
..
[c0487b]
Switch to side-by-side view
--- a +++ b/Inference/utils.py @@ -0,0 +1,157 @@ +import numpy as np +import scipy.signal +import scipy.io +from sklearn.preprocessing import StandardScaler +import pickle +from glob import glob +import os +import pyedflib +from tqdm import tqdm +from network import IncUNet_HR,IncUNet_BR + + +import pandas as pd +import torch +from torch.utils.data import TensorDataset,DataLoader +from torch.autograd import Variable + +def custom_resample(ECG,fs): + modified_ECG = [] + for i in range(int(len(ECG) * 500/fs)): + modified_ECG.append(ECG[int(fs/500*i)].astype(float)) + return modified_ECG + +def peak_correction(peak_locs,ecg_records): + + mod_peak_locs = [] + ecg_records = ecg_records.cpu().numpy() + for j in range(len(peak_locs)): + mod_peak_locs.append(np.asarray([peak_locs[j][i] - 37 + np.argmax(ecg_records[j,0,peak_locs[j][i]-37:peak_locs[j][i]+37]) for i in range(len(peak_locs[j])) if(peak_locs[j][i]>37 and peak_locs[j][i]<5000-37)])) + return mod_peak_locs + + +def obtain_ecg_record(path_records,fs,window_size = 5000): + + scaler = StandardScaler() + ecg_records = [] + actual_ecg_windows = [] + for path in path_records: + ECG = scipy.io.loadmat(path)['ecg'] + if(ECG.shape[0] != 1): + ECG = ECG[:,0] + else: + ECG = ECG[0,:] + ECG = np.asarray(custom_resample(ECG,fs)) + ### Scaling the ECG for every 5000 records(Speed vs Accuracy) + ecg_windows = [] + for record_no in range(len(ECG)//(window_size-500)): + ecg_windows.append(scaler.fit_transform(ECG[4500*record_no : 4500*record_no + window_size].reshape(-1,1))) + for record_no in range(len(ECG)//(window_size)): + actual_ecg_windows.append(scaler.fit_transform(ECG[5000*record_no : 5000*record_no + window_size].reshape(-1,1))) + ecg_records.append(np.asarray(ecg_windows)) + initial = np.zeros((1,5000)) + for ecg_record in ecg_records: + if(ecg_record[-1].shape[0] != 5000): + ecg_record[-1] = np.vstack((ecg_record[-1],np.zeros(5000 - ecg_record[-1].shape[0]).reshape(-1,1))) + for window_no in range(len(ecg_record)): + initial = np.vstack((initial,ecg_record[window_no][:,0].reshape(1,-1))) + + ecg_records = initial[1:,:] + actual_ecg_windows = np.asarray(actual_ecg_windows)[:,:,0] + return ecg_records,actual_ecg_windows + + return ecg_records + +def load_model_HR(SAVED_MODEL_PATH,test_loader,device,batch_len,window_size): + + C,H,W = 1,1,5000 + loaded_model = IncUNet_HR(in_shape=(C,H,W)) + loaded_model.load_state_dict(torch.load(SAVED_MODEL_PATH, map_location = lambda storage, loc: storage, pickle_module=pickle)) + loaded_model.to(device) + loaded_model.eval() + print("-------- Evaluation --------") + loaded_model.eval() + net_test_loss = 0 + pred_peaks = [] + peak_array = np.array([1]) + with torch.no_grad(): + pred_peak = [] + for step,(x) in tqdm(enumerate(test_loader)): + x = Variable(x[0].to(device)) + y_predict_test = loaded_model(x) + if(str(device)[:4] == 'cuda'): + y_predict_test = y_predict_test.cpu().numpy() + y_predict_test = y_predict_test[:,0,:] + else: + y_predict_test = y_predict_test.numpy() + y_predict_test = y_predict_test[:,0,:] + predicted_peak_locs = peak_finder(y_predict_test,x) + predicted_peak_locs = np.asarray(predicted_peak_locs) + for i in range(len(predicted_peak_locs)): + predicted_peak_locs_new = predicted_peak_locs[i][(predicted_peak_locs[i] >= 0.5*500) & (predicted_peak_locs[i] <= 9.5*500)] + pred_peaks.append(np.asarray(predicted_peak_locs_new + i*(window_size - 500) + step * batch_len * (window_size - 500)).astype(int)) + for i in range(len(pred_peaks)): + peak_array = np.hstack((peak_array,pred_peaks[i])) + actual_peak_locs = peak_array[1:] ### As peak locs were initialized with one. + return actual_peak_locs + +def load_model_BR(SAVED_MODEL_PATH,test_loader,device,batch_len,window_size): + + C,H,W = 1,1,5000 + loadedModel = IncUNet_BR(in_shape=(C,H,W)) + loadedModel.load_state_dict(torch.load(SAVED_MODEL_PATH)["model"]) + loadedModel.eval() + print("-------- Evaluation --------") + net_test_loss = 0 + + fs_upsample = 500 + fs_BR = 125 + no_sec = 10 + breathsPosition = np.array([]) + middle_start, middle_end = no_sec*fs_BR//4, 3*no_sec*fs_BR//4 + start = 0 + for step,x in enumerate(test_loader): + + valleys, predicted = np.array(getBR(x[0], loadedModel)) + current = valleys[(valleys >= middle_start) & ( valleys <= middle_end)] + startBR = int(start/(fs_upsample/fs_BR)) + if start == 0: + breathsPosition = np.append(breathsPosition,valleys[valleys < middle_end]) + else: + breathsPosition = np.append(breathsPosition, current + startBR) + + start += 2500 + return breathsPosition + +def peak_finder(y_pred_array,x_test): + + fs_ = 500 + peak_locs = [] + for i in range(y_pred_array.shape[0]): + peak_locs.append(scipy.signal.find_peaks(-y_pred_array[i,:],distance = 120,height = -0.4, prominence = 0.035)[0]) + peak_locs[i] = peak_locs[i][(peak_locs[i] >= 0.5*fs_) & (peak_locs[i] <= 9.5*fs_)] + modified_peak_locs = peak_correction(peak_locs,x_test) + return modified_peak_locs + +def compute_heart_rate(r_peaks): + fs_ = 500 + r_peaks = r_peaks[(r_peaks >= 0.5*fs_) & (r_peaks <= 9.5*fs_)] + return round( 60 * fs_ / np.mean(np.diff(r_peaks))) + +def smooth(signal,window_len=50): + y = pd.DataFrame(signal).rolling(window_len,center = True, min_periods = 1).mean().values.reshape((-1,)) + return y + +def findValleys(signal, prominence = 10, is_smooth = True , distance = 10): + """ Return prominent peaks and valleys based on scipy's find_peaks function """ + smoothened = smooth(-1*signal) + valley_loc = scipy.signal.find_peaks(smoothened, prominence=0.05, distance = 125)[0] + return valley_loc + +def getBR(signal, model): + model.eval() + with torch.no_grad(): + transformPredicted = model(signal) + transformPredicted = transformPredicted.cpu().numpy().reshape((-1,)) + valleys = findValleys(transformPredicted) + return valleys, transformPredicted \ No newline at end of file