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)