import os import scipy import numpy as np #!pip install wfdb import pandas as pd from mne.io import concatenate_raws, read_raw_edf import xml.etree.ElementTree as ET import multiprocessing BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'datasets') EPOCH_SEC_SIZE = 30 # ALL_CHANNELS = ['SaO2', 'H.R.', 'EEG(sec)', 'ECG', 'EMG', 'EOG(L)', 'EOG(R)', 'EEG', 'THOR RES', 'ABDO RES', 'POSITION', 'LIGHT', 'NEW AIR', 'OX stat'] AIRFLOW_NAMES = ["NEW AIR", "NEWAIR", "AIRFLOW", "AIRFLOW-0", "airflow", "new A/F", "New Air", "AUX"] EEG2_NAMES = ["EEG2", "EEG 2", "EEG(sec)", "EEG sec", "EEG(SEC)"] SELECT_CHNS = ['SaO2', 'ECG', 'EMG', 'EOG(L)', 'EOG(R)', 'EEG', 'ABDO RES'] + AIRFLOW_NAMES + EEG2_NAMES CHN_SAVE_DICT = { 'SaO2': ('SaO2','SaO2'), # 'EEG(sec)': ('EEG', 'C4-A1'), # 'EEG2': ('EEG', 'C4-A1'), # 'EEG 2': ('EEG', 'C4-A1'), # 'EEG sec': ('EEG', 'C4-A1'), 'ECG': ('ECG', 'ECG'), 'EMG': ('EMG', 'EMG'), 'EOG(L)': ('EOG', 'EOG-Diff'), 'EOG(R)': ('EOG', 'EOG-Diff'), 'EEG': ('EEG', 'C3-A2'), 'ABDO RES': ('ABD', 'ABD'), 'NEW AIR': ('AIRFLOW', 'AIRFLOW'), 'AIRFLOW': ('AIRFLOW', 'AIRFLOW') } import argparse parser = argparse.ArgumentParser() parser.add_argument('--data_dir', default="#_SHHS/polysomnography/edfs/shhs1", help='SHHS1 raw data location') parser.add_argument('--output_dir', default='shhs1/signals_MAE', help='SHHS1 processed data location') parser.add_argument('--annotation_dir', default='#_SHHS/polysomnography/annotations-events-profusion/shhs1', help='SHHS1 annotations location') parser.add_argument('--shhs2', action='store_true', help='Use SHHS2 dataset') cfg = parser.parse_args() from tqdm.auto import tqdm def process_signal(raw, select_idx, chns, fname): # sampling rate from instance raw_sampling_rate = int(raw.info['sfreq']) raw_df = raw.to_data_frame() # print(raw_df.head()) print(raw_df.columns) for chn in chns: if chn == "EOG(L)": raw_df['EOG-Diff'] = raw_df['EOG(L)'] - raw_df['EOG(R)'] raw_ch_df = raw_df['EOG-Diff'] elif chn == "EOG(R)": continue elif chn in AIRFLOW_NAMES: if chn in raw_df.columns: raw_ch_df = raw_df[chn] else: continue elif chn in EEG2_NAMES: if chn in raw_df.columns: raw_ch_df = raw_df[chn] else: continue else: raw_ch_df = raw_df[chn] raw_ch_df = raw_ch_df.to_frame() raw_ch_df.set_index(np.arange(len(raw_ch_df))) # Remove movement and unknown stages if any raw_ch = raw_ch_df.values # Verify that we can split into 30-s epochs if len(raw_ch) % (EPOCH_SEC_SIZE * raw_sampling_rate) != 0: print(f"Length of raw channel ({len(raw_ch)}) \ is not multiple of raw sampling rate and epoch size ({EPOCH_SEC_SIZE} * {raw_sampling_rate})") return np.array([]) # get (# of epochs) = (# of datapoints) / (30sec * 100Hz) n_epochs = len(raw_ch) / (EPOCH_SEC_SIZE * raw_sampling_rate) # Get epochs and their corresponding labels x = np.asarray(np.split(raw_ch, n_epochs)).astype(np.float32) # split epochs x = x[select_idx] x = x.reshape((-1, EPOCH_SEC_SIZE * raw_sampling_rate)) if chn in AIRFLOW_NAMES: chn_grp, chn_rename = 'AIRFLOW', 'AIRFLOW' elif chn in EEG2_NAMES: chn_grp, chn_rename = 'EEG', 'C4-A1' else: chn_grp, chn_rename = CHN_SAVE_DICT[chn] # chn_grp = CHN_SAVE_DICT[chn][0] if chn not in AIRFLOW_NAMES else 'AIRFLOW' # chn_rename = CHN_SAVE_DICT[chn][1] if chn not in AIRFLOW_NAMES else 'AIRFLOW' # print(chn_grp, chn_rename) os.makedirs(os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn_rename), exist_ok=True) np.save(os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn_rename, f"{fname}.npy"), x) def process_data(args): ann_fname, edf_fname, file_id, shhs2 = args print(args) ann_filename = os.path.basename(edf_fname).replace(".edf", "_ann.npy") # if os.path.exists(os.path.join(BASE_PATH, cfg.output_dir, "annotations", ann_filename)): # return labels = [] # Read annotation and its header t = ET.parse(ann_fname) r = t.getroot() faulty_File = 0 for i in range(len(r[4])): lbl = int(r[4][i].text) # print(lbl) if lbl == 4: # make stages N3, N4 same as N3 labels.append(3) elif lbl == 5: # Assign label 4 for REM stage labels.append(4) else: labels.append(lbl) if lbl > 5: # some files may contain labels > 5 BUT not the selected ones. faulty_File = 1 if faulty_File == 1: print("Faulty labels found!!") return -1 labels = np.asarray(labels) # label list to npaprray # corresponding labels y = labels.astype(np.int32) # Select on sleep periods w_edge_mins = 30 # get index of first not WAKE nw_idx = np.where(y != 0)[0] # start from 30min before first WAKE start_idx = nw_idx[0] - (w_edge_mins * 2) # end at 30min after first WAKE end_idx = nw_idx[-1] + (w_edge_mins * 2) if start_idx < 0: start_idx = 0 if end_idx >= len(y): end_idx = len(y) - 1 select_idx = np.arange(start_idx, end_idx + 1) # only selected labels y = y[select_idx] # get RawEDF instance # !!!!!!!!!!!!!!!!!! 12~13 sec raw = read_raw_edf(edf_fname, preload=True, stim_channel=None, verbose=None) process_signal(raw, select_idx, SELECT_CHNS, os.path.basename(edf_fname).split('.')[0]) ann_dir = os.path.join(BASE_PATH, 'shhs1' if not shhs2 else 'shhs2', "annotations") os.makedirs(ann_dir, exist_ok=True) np.save(os.path.join(ann_dir, ann_filename), arr=y) 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)) file_list = [f for f in file_list if f.startswith('tr')] filenames_edf = pd.read_csv("shhs_meta_data/filenames_edf.txt", header=None, names=['a']) filenames_edf = filenames_edf['a'].values.tolist() filenames_ann = pd.read_csv("shhs_meta_data/filenames_ann.txt", header=None, names=['a']) filenames_ann = filenames_ann['a'].values.tolist() if cfg.shhs2: edf_fnames = [os.path.join(BASE_PATH, cfg.data_dir, f) for f in os.listdir(os.path.join(BASE_PATH, cfg.data_dir))] ann_fnames = [os.path.join(BASE_PATH, cfg.annotation_dir, f) for f in os.listdir(os.path.join(BASE_PATH, cfg.annotation_dir))] else: # make actual file names for selected edf files and make a list edf_fnames = [os.path.join(BASE_PATH, cfg.data_dir, filename_edf) for filename_edf in filenames_edf] ann_fnames = [os.path.join(BASE_PATH, cfg.annotation_dir, filename_ann) for filename_ann in filenames_ann] edf_fnames.sort() # sort edf file names list ann_fnames.sort() # sort ann file names list # sorted edf/ann file names list to nparray # edf_fnames = np.asarray(edf_fnames) # ann_fnames = np.asarray(ann_fnames) input_list = zip(ann_fnames, edf_fnames, range(len(ann_fnames)), len(ann_fnames) * [cfg.shhs2]) # redo_list = ['shhs1-203353.npy', 'shhs1-203171.npy', 'shhs1-203301.npy', 'shhs1-203011.npy', 'shhs1-203300.npy', 'shhs1-203419.npy', 'shhs1-203286.npy', 'shhs1-203375.npy','shhs1-201820.npy'] # input_list = [args for args in input_list if os.path.basename(args[1]).replace('.edf', '.npy') in redo_list] # print(len(input_list)) # exit(0) # for args in input_list: # ann_fname, edf_fname, file_id = args # if 'shhs1-201820' in edf_fname: # process_data(args) # break with multiprocessing.Pool(8) as pool: list(tqdm(pool.imap(process_data, input_list), total=len(ann_fnames)))