Download data/cfs_process_rawdata.py from gundo0102/SleepMaMi: direct link, hf CLI and curl.
- Browser
- Download file 11.1 kB
-
https://huggingface.co/gundo0102/SleepMaMi/resolve/main/data/cfs_process_rawdata.py
- Command line
-
hf download hf://gundo0102/SleepMaMi/data/cfs_process_rawdata.py
-
curl -L -o cfs_process_rawdata.py https://huggingface.co/gundo0102/SleepMaMi/resolve/main/data/cfs_process_rawdata.py
11.1 kB
| 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))) | |