100 lines
3.9 KiB
Python
100 lines
3.9 KiB
Python
import pandas as pd
|
||
import numpy as np
|
||
import matplotlib
|
||
import matplotlib.pyplot as plt
|
||
import ecgdetectors
|
||
from ecgdetectors import Detectors
|
||
from scipy import fftpack
|
||
import os
|
||
from BCGDataset import BCGDataset,BCG_Operation,read_all_data
|
||
import math
|
||
|
||
def refinement( data, peak):
|
||
if len(data) == 0 or len(peak) <=2 : return None
|
||
firstPeak = peak[0]
|
||
lastPeak = peak[-1]
|
||
meanPeak = np.quantile( data[peak[1:-1]], 0.2 )
|
||
if data[firstPeak] < meanPeak * 0.6 :
|
||
peak = np.delete(peak, 0)
|
||
if data[lastPeak] < meanPeak * 0.6 :
|
||
peak = np.delete(peak, -1)
|
||
return np.array(peak)
|
||
|
||
def find_TPeak(data,peaks,th=50):
|
||
"""
|
||
找出真实的J峰或R峰
|
||
:param data: BCG或ECG数据
|
||
:param peaks: 初步峰值（从label中导出的location_R）
|
||
:param th: 范围阈值
|
||
:return: 真实峰值
|
||
"""
|
||
return_peak = []
|
||
for peak in peaks:
|
||
if peak>len(data):continue
|
||
min_win,max_win = max(0,int(peak-th)),min(len(data),int(peak+th))
|
||
return_peak.append(np.argmax(data[min_win:max_win])+min_win)
|
||
return np.array(return_peak)
|
||
|
||
def Rpeak_Detection(raw_ecg,fs,low_cut,high_cut,th1,detector_method):
|
||
detectors = Detectors(sampling_frequency=fs)
|
||
method_dic = {'pt': detectors.pan_tompkins_detector,
|
||
'ta': detectors.two_average_detector,
|
||
"Engzee": detectors.engzee_detector,
|
||
"Wt": detectors.swt_detector,
|
||
"Christov": detectors.christov_detector,
|
||
"Hamilton": detectors.hamilton_detector
|
||
}
|
||
detectormethods = method_dic[detector_method]
|
||
|
||
# raw_ecg = raw_ecg[200*sample_rate:]
|
||
preprocessing = BCG_Operation(sample_rate=fs) # 对ECG做了降采样处理
|
||
raw_ecg = preprocessing.Butterworth(raw_ecg, "bandpass", low_cut=low_cut, high_cut=high_cut, order=3) * 4
|
||
#######################限制幅值处理############################################
|
||
# for i in range(len(raw_ecg)):
|
||
# if raw_ecg[i] > 300 or raw_ecg[i] < -300:
|
||
# raw_ecg[i] = 0
|
||
##############################################################################
|
||
|
||
##############################切割处理##########################################
|
||
all_file_num = math.ceil((len(raw_ecg) / fs / 60 / 60))
|
||
ecg_seq = np.array(np.arange(all_file_num))
|
||
ecg_seq = ecg_seq.astype(np.ndarray)
|
||
R_peak_seq = np.array(np.arange(all_file_num))
|
||
R_peak_seq = R_peak_seq.astype(np.ndarray)
|
||
Interval_seq = np.array(np.arange(all_file_num))
|
||
Interval_seq = Interval_seq.astype(np.ndarray)
|
||
RRIV_seq = np.array(np.arange(all_file_num))
|
||
RRIV_seq = RRIV_seq.astype(np.ndarray)
|
||
|
||
for file_num in range(1,all_file_num+1):
|
||
if file_num != all_file_num:
|
||
new_ecg = raw_ecg[fs*3600*(file_num - 1) : fs*3600*file_num]
|
||
else:
|
||
new_ecg = raw_ecg[fs*3600*(file_num - 1):]
|
||
##############################################################################
|
||
|
||
R_peak = np.array(detectormethods(new_ecg)) - 100
|
||
# R_peak = np.array(detectors.pan_tompkins_detector(raw_ecg))-100
|
||
|
||
R_peak = find_TPeak(new_ecg, R_peak, th=int(th1 * fs / 1000))
|
||
R_peak = refinement(new_ecg, R_peak)
|
||
|
||
RR_Interval = np.full(len(R_peak) - 1, np.nan)
|
||
|
||
for i in range(len(R_peak) - 1):
|
||
RR_Interval[i] = R_peak[i + 1] - R_peak[i]
|
||
|
||
RRIV = np.full(len(RR_Interval) - 1, np.nan)
|
||
for i in range(len(RR_Interval) - 1):
|
||
RRIV[i] = RR_Interval[i + 1] - RR_Interval[i]
|
||
|
||
Interval = np.full(len(new_ecg), np.nan)
|
||
for i in range(len(R_peak) - 1):
|
||
Interval[R_peak[i]: R_peak[i + 1]] = R_peak[i + 1] - R_peak[i]
|
||
|
||
ecg_seq[file_num - 1] = new_ecg
|
||
R_peak_seq[file_num - 1] = R_peak
|
||
Interval_seq[file_num - 1] = Interval
|
||
RRIV_seq[file_num - 1] = RRIV
|
||
|
||
return ecg_seq, R_peak_seq, Interval_seq, RRIV_seq |