智能诊断实验6

代码：

import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import butter, filtfilt, medfilt
from scipy.fftpack import fft
import scipy.io as scio
plt.rcParams['font.sans-serif'] = ['SimHei']
dataFile_100 = r"C:\Users\塔伯\Desktop\sample2017\sample2017\validation\A00001.mat"
data_100 = scio.loadmat(dataFile_100)
for val in data_100.keys():
    Y_100 = data_100[val]

NECG1 = Y_100[0][:1000]
# 基线去除
baseline_removed_signal1 = NECG1 - medfilt(NECG1, kernel_size=1001)
fs = 1000
fc = 50
Rp = 3
Rs = 60
Wp = 100 / (fs / 2)
Ws = 120 / (fs / 2)

# 计算滤波器阶数
n = np.log10((10**(0.1*Rs)-1) / (10**(0.1*Rp)-1)) / (2 * np.log10(Ws/Wp))

# 计算截止频率
Wn = fc / (fs / 2)
b, a = butter(4, fc / (fs / 2), 'low')

# 进行滤波
filtered_signal1 = filtfilt(b, a, baseline_removed_signal1)

# 进行傅立叶变换
X1 = fft(NECG1)
f = np.fft.fftfreq(len(X1), 1/fs)

# 滤波前的信号
x1_filtered = filtfilt(b, a, NECG1)
t1 = np.arange(0, len(NECG1)) / fs
n1 = NECG1 - x1_filtered
SNR1 = 10 * np.log10(np.sum(NECG1**2) / np.sum(n1**2))

# 滤波后的信号
n2 = NECG1 - filtered_signal1
SNR2 = 10 * np.log10(np.sum(NECG1**2) / np.sum(n2**2))

plt.figure(figsize=(10, 8))

plt.subplot(2, 2, 1)
plt.plot(f[:len(X1)//2], np.abs(X1[:len(X1)//2]))
plt.title('原始心电信号频谱图')
plt.xlabel('频率（Hz）')
plt.ylabel('幅值')

plt.subplot(2, 2, 2)
plt.plot(t1, NECG1)
plt.title('原始心电信号时域图')
plt.xlabel('时间（s）')
plt.ylabel('幅值')

X1_filtered = fft(x1_filtered)
plt.subplot(2, 2, 3)
plt.plot(f[:len(X1)//2], np.abs(X1_filtered[:len(X1)//2]))
plt.title('滤波后心电信号频谱图')
plt.xlabel('频率（Hz）')
plt.ylabel('幅值')

plt.subplot(2, 2, 4)
plt.plot(t1, filtered_signal1)
plt.title('增强后的心电信号时域图')
plt.xlabel('时间（s）')
plt.ylabel('幅值')

plt.tight_layout()
plt.show()

print("截止频率：", fc)
print("增强后的信号的 SNR：", SNR1)
print("原始信号的 SNR：", SNR2)