test_a_sig.py

# -*- coding: utf-8 -*-
"""
Created on Sun Apr 21 14:08:55 2019

@author: Winham

test_a_sig.py: 加载训练好的模型，从验证集中随机选取一条信号进行测试

"""

import os
import numpy as np
import tensorflow as tf
from keras.models import load_model
import keras.backend as K
from sklearn import preprocessing as prep
import matplotlib.pyplot as plt
import time
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

val_sig_path = 'G:/ECG_UNet/val_sigs/'
val_label_path = 'G:/ECG_UNet/val_labels/'

sig_files = os.listdir(val_sig_path)
label_files = os.listdir(val_label_path)

select = np.random.choice(sig_files, 1)[0]

a_sig = np.load(val_sig_path+select)
a_seg = np.load(val_label_path+select)

K.clear_session()
tf.reset_default_graph()
model = load_model('myNet.h5')

a_sig = np.expand_dims(prep.scale(a_sig), axis=1)
a_sig = np.expand_dims(a_sig, axis=0)

tic = time.time()
a_pred = model.predict(a_sig)
toc = time.time()


print('Elapsed time: '+str(toc-tic)+' seconds.')
plt.plot(a_sig[0, :, 0])
plt.grid(True)

plt.plot(a_pred[0, :, 0], 'b')
plt.plot(a_pred[0, :, 1], 'k')
plt.plot(a_pred[0, :, 2], 'r')
plt.legend(['Sig', 'Background', 'Normal', 'PVC'], loc='lower right')