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') # BASE_PATH = 'datasets' EPOCH_SEC_SIZE = 30 # 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', 'SpO2','ECG1', 'EMG1', 'LOC', 'ROC', 'C3', 'C4', 'M1', 'M2', 'A1','A2', 'ABDO EFFORT', 'THOR EFFORT', 'AIRFLOW'] CHN_SAVE_DICT = { 'SaO2': ('SaO2','SaO2'), 'SpO2': ('SaO2','SaO2'), 'ECG1': ('ECG', 'ECG'), 'EMG1': ('EMG', 'EMG'), 'LOC': ('EOG', 'EOG-Diff'), 'ROC': ('EOG', 'EOG-Diff'), 'C3': ('EEG', 'C3-M2'), 'C4': ('EEG', 'C4-M1'), 'M1': ('EEG', 'C4-M1'), 'M2': ('EEG', 'C3-M2'), 'A1': ('EEG', 'A1-M2'), 'A2': ('EEG', 'A2-M1'), 'ABDO EFFORT': ('ABD', 'ABD'), 'THOR EFFORT': ('CHEST', 'CHEST'), 'AIRFLOW': ('AIRFLOW', 'AIRFLOW') } import argparse parser = argparse.ArgumentParser() parser.add_argument('--data_dir', default="CFS_raw/polysomnography/edfs", help='CFS raw data location') parser.add_argument('--output_dir', default='cfs/signals_MAE', help='CFS processed data location') parser.add_argument('--annotation_dir', default='CFS_raw/polysomnography/annotations-events-profusion', help='CFS annotations location') parser.add_argument('--ah_dir', default='CFS_raw/polysomnography/annotations-events-nsrr', help='CFS apnea/hypopnea labels location') parser.add_argument('--demographics_file', default='CFS_raw/datasets/cfs-visit5-harmonized-dataset-0.7.0.csv', help='CFS demographics file location') 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 == "LOC": raw_df['EOG-Diff'] = raw_df['LOC'] - raw_df['ROC'] raw_ch_df = raw_df['EOG-Diff'] elif chn == "ROC": continue elif chn == 'C3': if 'M2' in raw_df.columns: raw_df['C3-M2'] = raw_df['C3'] - raw_df['M2'] else: raw_df['C3-M2'] = raw_df['C3'] - raw_df['A2'] raw_ch_df = raw_df['C3-M2'] elif chn == 'C4': if 'M1' in raw_df.columns: raw_df['C4-M1'] = raw_df['C4'] - raw_df['M1'] else: raw_df['C4-M1'] = raw_df['C4'] - raw_df['A1'] raw_ch_df = raw_df['C4-M1'] elif chn in ['M1', 'M2', 'A1', 'A2']: continue elif chn in ['SaO2', 'SpO2']: if 'SaO2' in raw_df.columns: raw_ch_df = raw_df['SaO2'] else: raw_ch_df = raw_df['SpO2'] 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)) 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, ah_fname, edf_fname, file_id = 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, 'cfs', 'annotations') os.makedirs(ann_dir, exist_ok=True) np.save(os.path.join(ann_dir, ann_filename), arr=y) tree = ET.parse(ah_fname) root = tree.getroot() events = root.findall('.//ScoredEvent') total_seconds = 0 for event in events: if event.find('EventConcept').text == 'Recording Start Time': duration = float(event.find('Duration').text) total_seconds = int(np.ceil(duration)) break # 만약 Recording Start Time을 못 찾았을 경우 (예외 처리), 기존 방식대로 마지막 이벤트 시간으로 계산 if total_seconds == 0: print("예외처리로 인한 Start time ~ End time 계산") max_time = 0 for event in events: try: start = float(event.find('Start').text) dur = float(event.find('Duration').text) if start + dur > max_time: max_time = start + dur except: continue total_seconds = int(np.ceil(max_time)) # 2. 라벨 어레이 생성 ah_labels = np.zeros(total_seconds, dtype=int) # 3. 호흡 이벤트 파싱 및 라벨링 for event in events: try: event_type = event.find('EventType').text if event_type != 'Respiratory|Respiratory': continue # 호흡 이벤트가 아니면 스킵 concept = event.find('EventConcept').text start = float(event.find('Start').text) duration = float(event.find('Duration').text) start_idx = int(round(start)) end_idx = start_idx + int(round(duration)) # 인덱스 범위 보정 if end_idx > total_seconds: end_idx = total_seconds # 라벨링 (Hypopnea=1, Apnea=2) if 'Hypopnea' in concept: mask = ah_labels[start_idx:end_idx] != 2 ah_labels[start_idx:end_idx][mask] = 1 elif any(x in concept for x in ['Obstructive', 'Central', 'Mixed']): ah_labels[start_idx:end_idx] = 2 except AttributeError: continue ah_labels = np.reshape(ah_labels, (-1, 30)) ah_labels = ah_labels[select_idx] assert len(ah_labels) == len(y), f"Length of ah_labels ({len(ah_labels)}) does not match length of sleep stage labels ({len(y)})" ah_labels = np.reshape(ah_labels, (-1,)) ah_dir = os.path.join(BASE_PATH, 'cfs', 'ah_labels') ah_filename = os.path.basename(edf_fname).replace(".edf", "_ah.npy") os.makedirs(ah_dir, exist_ok=True) np.save(os.path.join(ah_dir, ah_filename), arr=ah_labels) 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)) demo = pd.read_csv(os.path.join(BASE_PATH, cfg.demographics_file)) def get_nsrrid(filename): nsrrid = filename.split('-')[-1].replace('.edf', '') return int(nsrrid) age_filtered_df = demo[demo['nsrr_age'] > 13] file_list = [filename for filename in file_list if get_nsrrid(filename) in age_filtered_df['nsrrid'].values] print(f"Number of files after filtering by age > 13: {len(file_list)}") filenames_edf = [filename for filename in file_list if filename.endswith('.edf')] filenames_ann = [filename.replace('.edf', '-profusion.xml') for filename in filenames_edf] filenames_ah = [filename.replace('.edf', '-nsrr.xml') for filename in filenames_edf] 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] ah_fnames = [os.path.join(BASE_PATH, cfg.ah_dir, filename_ah) for filename_ah in filenames_ah] edf_fnames.sort() # sort edf file names list ann_fnames.sort() # sort ann file names list ah_fnames.sort() # sort ah file names list input_list = zip(ann_fnames, ah_fnames, edf_fnames, range(len(ann_fnames))) # 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)))