File size: 7,198 Bytes
35f60b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | import os
import scipy
import numpy as np
#!pip install wfdb
import wfdb
import math
BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'datasets')
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--data_dir', default='physionet.org/files/challenge-2018/1.0.0/training/', help='Physionet 2018 raw data location')
parser.add_argument('--output_dir', default='physionet_2018/signals_MAE', help='Physionet 2018 processed data location')
parser.add_argument('--annotation_dir', default='physionet_2018/annotations', help='Physionet 2018 annotations location')
parser.add_argument('--ah_dir', default='physionet_2018/ah_labels', help='Physionet 2018 apnea/hypopnea labels location')
parser.add_argument('--channels', nargs='+', default=[0,1,2,3,4,5,6,7,8,9,10,11,12], help='Physionet 2018 annotations location')
parser.add_argument('--test', action='store_true', help='Preprocess test set')
cfg = parser.parse_args()
from tqdm.auto import tqdm
if __name__ == '__main__':
# data_dir = '/tf/00_data/kdpark/physionet_2018/physionet_rawdata'
file_list = os.listdir(os.path.join(BASE_PATH, cfg.data_dir))
if cfg.test:
file_list = [f for f in file_list if f.startswith('te')]
else:
file_list = [f for f in file_list if f.startswith('tr')]
os.makedirs(os.path.join(BASE_PATH, cfg.output_dir), exist_ok=True)
os.makedirs(os.path.join(BASE_PATH, cfg.annotation_dir), exist_ok=True)
os.makedirs(os.path.join(BASE_PATH, cfg.ah_dir), exist_ok=True)
def match_channel_group(chn):
"""
Find the group of the channel.
"""
chn_group = {
'EEG': ['F3-M2', 'F4-M1', 'C3-M2', 'C4-M1', 'O1-M2', 'O2-M1'],
'EOG': ['E1-M2'],
'EMG': ['Chin1-Chin2'],
'ABD': ['ABD'],
'CHEST': ['CHEST'],
'AIRFLOW': ['AIRFLOW'],
'SaO2': ['SaO2'],
'ECG': ['ECG']
}
for k, v in chn_group.items():
if chn in v:
return k
return None
for fl in tqdm(file_list):
# EEG: [C4-M1], EOG: [E1-M2], EMG: [Chin1-Chin2]
# fields: {'fs': 200, 'sig_len': 5569600, 'n_sig': 13, 'base_date': None,
# 'base_time': None, 'units': ['uV', 'uV', 'uV', 'uV', 'uV', 'uV', 'uV', 'uV', 'uV', 'uV', 'uV', '%', 'mV'],
# 'sig_name': ['F3-M2', 'F4-M1', 'C3-M2', 'C4-M1', 'O1-M2', 'O2-M1', 'E1-M2', 'Chin1-Chin2', 'ABD', 'CHEST', 'AIRFLOW', 'SaO2', 'ECG'], 'comments': []}
if fl == 'tr07-0023': continue
signals, fields = wfdb.rdsamp(os.path.join(BASE_PATH, cfg.data_dir, fl, fl), channels=cfg.channels)
if cfg.test:
num_epoch = signals.shape[0] // (200*30)
mat_signal = signals.astype(np.float32) #Shape: (num_samples, 13)
mat_signal = mat_signal[:num_epoch*200*30,:]
else:
ann = wfdb.rdann(os.path.join(BASE_PATH, cfg.data_dir, fl, fl),extension='arousal')
an = np.array(ann.aux_note) #Sleep stage and arousals annotations
s_time = ann.sample #annotation locations in samples relative to the beginning of the record.
an_where = np.concatenate((np.where(an=='W'),np.where(an=='N1'),np.where(an=='N2'),np.where(an=='N3'),np.where(an=='R')), axis=-1)
an_where.sort()
ah_where = np.concatenate((np.where(an=='(resp_obstructiveapnea'),
np.where(an=='resp_obstructiveapnea)'),
np.where(an=='(resp_hypopnea'),
np.where(an=='resp_hypopnea)')), axis=-1)
ah_where.sort()
stages = an[an_where[0]]
stages_time = s_time[an_where[0]]
ah = an[ah_where[0]]
ah_time = s_time[ah_where[0]]
print(len(an), len(stages_time))
num_epoch = int(stages_time[-1]/200/30) - int(stages_time[0]/200/30)
labels = -1 * np.ones((num_epoch,))
ah_labels = np.zeros((num_epoch*30,))
curr_epoch = 0
status= 0 if stages[0]=='W' else 1 if stages[0]=='N1' else 2 if stages[0]=='N2' else 3 if stages[0]=='N3' else 4
for i in range(1, len(stages)):
labels[curr_epoch: int(stages_time[i]/200/30) - int(stages_time[0]/200/30)] = status
curr_epoch = int(stages_time[i]/200/30 - int(stages_time[0]/200/30))
if stages[i] == 'W':
status = 0
elif stages[i] == 'N1':
status = 1
elif stages[i] == 'N2':
status = 2
elif stages[i] == 'N3':
status = 3
elif stages[i] == 'R':
status = 4
else:
status = -1
continue
curr_ah_label = None
for ah_label, ah_time_point in zip(ah, ah_time):
if ah_label == '(resp_obstructiveapnea' or ah_label == '(resp_hypopnea':
start = ah_time_point
curr_ah_label = ah_label
elif ah_label == 'resp_obstructiveapnea)' or ah_label == 'resp_hypopnea)':
end = ah_time_point
if curr_ah_label is None or curr_ah_label[1:] != ah_label[:-1]:
print(f"file:{fl}, Error in apnea/hypopnea annotation!, start_ah_label:{curr_ah_label}, end_ah_label:{ah_label}, from:{start}, to:{end}")
continue
start_sec = round((start - stages_time[0]) / 200)
end_sec = round((end - stages_time[0]) / 200)
end_sec = min(end_sec, len(ah_labels))
if ah_label == 'resp_obstructiveapnea)':
ah_labels[start_sec:end_sec] = 2
else:
ah_labels[start_sec:end_sec] = 1
# print(start_sec, end_sec, ah_label)
#Start time of signal&ann are always the same. Only match the end time
mat_signal = signals[stages_time[0]:stages_time[-1]].astype(np.float32) #Shape: (num_samples, 13)
if status == -1:
print("-1 error!")
print(fl, stages_time[0], stages_time[-1], len(labels)*200*30, len(mat_signal))
for i, fld in enumerate(fields['sig_name']):
chn_grp = match_channel_group(fld)
os.makedirs(os.path.join(BASE_PATH, cfg.output_dir, chn_grp, fld), exist_ok=True)
np.save(os.path.join(BASE_PATH, cfg.output_dir, chn_grp, fld, fl + '.npy'), mat_signal[:,i].reshape([num_epoch, 200*30]))
if not cfg.test:
# np.save(os.path.join(BASE_PATH, cfg.output_dir, fl + '.npy'), mat_signal)
np.save(os.path.join(BASE_PATH, cfg.annotation_dir, fl + '_ann.npy'), labels)
np.save(os.path.join(BASE_PATH, cfg.ah_dir, fl + '_ah.npy'), ah_labels)
|