| import os |
| import scipy |
| import numpy as np |
| |
| 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__': |
|
|
| |
| 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): |
|
|
| |
| |
| |
| |
|
|
| 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) |
| 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) |
| s_time = ann.sample |
|
|
| 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 |
| |
| |
|
|
| |
| mat_signal = signals[stages_time[0]:stages_time[-1]].astype(np.float32) |
|
|
| 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.annotation_dir, fl + '_ann.npy'), labels) |
| np.save(os.path.join(BASE_PATH, cfg.ah_dir, fl + '_ah.npy'), ah_labels) |
|
|
|
|
