import scipy.io import os from tqdm import tqdm import numpy as np import mne import shutil CHN_SAVE_DICT = { 'C3-A2': ('EEG', 'C3-A2'), 'C4-A1': ('EEG', 'C4-A1'), 'F3-A2': ('EEG', 'F3-A2'), 'F4-A1': ('EEG', 'F4-A1'), 'O1-A2': ('EEG', 'O1-A2'), 'O2-A1': ('EEG', 'O2-A1'), 'C3-M2': ('EEG', 'C3-A2'), 'C4-M1': ('EEG', 'C4-A1'), 'F3-M2': ('EEG', 'F3-A2'), 'F4-M1': ('EEG', 'F4-A1'), 'O1-M2': ('EEG', 'O1-A2'), 'O2-M1': ('EEG', 'O2-A1'), 'LOC-A2': ('EOG', 'Left-A2'), 'ROC-A1': ('EOG', 'Right-A1'), 'E1-M2': ('EOG', 'Left-A2'), 'E2-M1': ('EOG', 'Right-A1'), 'X1': ('EMG', 'EMG'), 'X2': ('ECG', 'ECG'), 'X6': ('AIRFLOW', 'AIRFLOW'), 'X7': ('CHEST', 'CHEST'), 'X8': ('ABD', 'ABD'), '24': ('EMG', 'EMG'), '25': ('ECG', 'ECG'), '29': ('AIRFLOW', 'AIRFLOW'), '30': ('CHEST', 'CHEST'), '31': ('ABD', 'ABD'), 'SpO2': ('SaO2', 'SaO2'), 'SaO2': ('SaO2', 'SaO2') } SELECT_CHNS1 = ['F3-A2', 'F4-A1', 'C3-A2', 'C4-A1', 'O1-A2', 'O2-A1', 'LOC-A2', 'ROC-A1', 'X1', 'X2', 'X6', 'X7', 'X8', 'SaO2'] SELECT_CHNS2 = ['F3-M2', 'F4-M1', 'C3-M2', 'C4-M1', 'O1-M2', 'O2-M1', 'E1-M2', 'E2-M1', '24', '25', '29', '30', '31', 'SpO2'] # def match_channel_group(chn): # """ # Find the group of the channel. # """ # chn_group = { # 'EEG': ['F3-A2', 'F4-A1', 'C3-A2', 'C4-A1', 'O1-A2', 'O2-A1'], # 'EOG': ['LOC-A2', 'ROC-A1'], # 'EMG': ['X1'], # 'ABD': ['X8'], # 'CHEST': ['X7'], # 'AIRFLOW': ['X6'], # 'SaO2': ['SaO2'], # 'ECG': ['X2'] # } # for k, v in chn_group.items(): # if chn in v: # return k # return None BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), 'datasets') import argparse parser = argparse.ArgumentParser() parser.add_argument('--data_dir', default='ISRUC', help='ISRUC raw data location') parser.add_argument('--output_dir', default='ISRUC', help='ISRUC processed data location') parser.add_argument('--annotation_dir', default='ISRUC', help='ISRUC annotations location') cfg = parser.parse_args() def read_ann(ann_path): # Annotation with open(ann_path, 'r') as f: ann = f.readlines() ann = np.array([line.strip() for line in ann if line.strip()]).astype(np.int64) ann = np.where(ann==5, 4, ann) return ann # Example usage raw_sampling_rate = 200 EPOCH_SEC_SIZE = 30 for ds_id in [1,2,3]: if ds_id == 1: num_subject = 100 elif ds_id == 2: num_subject = 8 else: num_subject = 10 num_session = 2 if ds_id == 2 else 1 for i in tqdm(range(1, num_subject+1)): for j in range(1, num_session+1): data_dir_full = os.path.join(BASE_PATH, cfg.data_dir, f'SG{ds_id}', f'{i}',f'{i}') if ds_id == 2: file_path = os.path.join(data_dir_full, f'{j}', f'{j}.edf') ann_path = os.path.join(data_dir_full, f'{j}', f'{j}_1.txt') if not os.path.exists(file_path): rarpath = os.path.join(data_dir_full, f'{j}', f'{j}.rec') shutil.copy(rarpath, file_path) else: file_path = os.path.join(data_dir_full,f'{i}.edf') ann_path = os.path.join(data_dir_full,f'{i}_1.txt') if not os.path.exists(file_path): rarpath = os.path.join(data_dir_full, f'{i}.rec') shutil.copy(rarpath, file_path) ann = read_ann(ann_path) raw = mne.io.read_raw_edf(file_path, preload=True) raw_df = raw.to_data_frame() if not len(raw_df) == len(ann) * EPOCH_SEC_SIZE * raw_sampling_rate: print(f"Warning: {i} signal length and annotation length mismatched") # Try to use the second annotator's scoring if ds_id == 2: ann_path = os.path.join(data_dir_full, f'{j}', f'{j}_2.txt') else: ann_path = os.path.join(data_dir_full,f'{i}_2.txt') ann = read_ann(ann_path) if not len(raw_df) == len(ann) * EPOCH_SEC_SIZE * raw_sampling_rate: print(f"Error both annotations mismatch signal length, skip this record. " , "Dataset id:", ds_id, "Subject id:", i) continue ann_dir = os.path.join(BASE_PATH, f'{cfg.output_dir}_SG{ds_id}', 'annotations') os.makedirs(ann_dir, exist_ok=True) fname = f'{i:03}_{j}' if ds_id == 2 else f'{i:03}' ann_name = f'{i:03}_{j}_ann' if ds_id == 2 else f'{i:03}_ann' np.save(os.path.join(ann_dir, ann_name), ann) for chn in SELECT_CHNS1: # chn_grp = match_channel_group(chn) if chn in raw_df.columns: raw_ch_df = raw_df[chn] chn_grp, chn_rename = CHN_SAVE_DICT[chn] else: idx = SELECT_CHNS1.index(chn) chn = SELECT_CHNS2[idx] if chn in raw_df.columns: raw_ch_df = raw_df[chn] chn_grp, chn_rename = CHN_SAVE_DICT[chn] else: print(f"Channel {chn} not found in {file_path}, skip this channel.") print(raw_df.columns) continue 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_df) % (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})") # 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)) os.makedirs(os.path.join(BASE_PATH, f'{cfg.output_dir}_SG{ds_id}', 'signals_MAE', chn_grp, chn_rename), exist_ok=True) np.save(os.path.join(BASE_PATH, f'{cfg.output_dir}_SG{ds_id}', 'signals_MAE', chn_grp, chn_rename, f"{fname}.npy"), x)