| import os |
| import shutil |
| import numpy as np |
| import matplotlib.pyplot as plt |
| from tqdm import tqdm |
| import shutil |
| import librosa |
|
|
| from scipy.signal import iirnotch, lfilter |
| from data_utils import butter_bandpass_filter, correct_ecg_inversion |
| import pickle |
| import pandas as pd |
|
|
| BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| NOTCH_FREQ = 60 |
| EEG_LOWCUT = 0.3 |
| EEG_HIGHCUT = 40 |
| EMG_LOWCUT = 10 |
| EMG_HIGHCUT = 1000 |
| ECG_LOWCUT = 0.3 |
| ECG_HIGHCUT = 60 |
| FLOW_LOWCUT = 0 |
| FLOW_HIGHCUT = 15 |
|
|
|
|
|
|
| NUM_PROCESS = 4 |
|
|
| import argparse |
| import multiprocessing |
|
|
|
|
|
|
| |
| TARGET_RESAMPLE = 100 |
|
|
| |
|
|
|
|
|
|
|
|
|
|
| def process_ecg(args): |
| sig, chn_grp, chn = args |
| sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig) |
| data = np.load(sig_path) |
| sampling_rate = data.shape[1] // (30) |
| |
| data = data.reshape(-1, 1) |
| |
| filtered = butter_bandpass_filter(data, ECG_LOWCUT, ECG_HIGHCUT, sampling_rate) |
| b_notch, a_notch = iirnotch(NOTCH_FREQ, 30, sampling_rate) |
| filtered = lfilter(b_notch, a_notch, filtered) |
|
|
| filtered = filtered.reshape(-1, sampling_rate*30) |
| |
| |
| filtered_epochs = [] |
| |
| for epoch in filtered: |
| inv_epoch, _ = correct_ecg_inversion(epoch, TARGET_RESAMPLE) |
| filtered_epochs.append(inv_epoch) |
|
|
| filtered = np.stack(filtered_epochs, axis=0) |
|
|
| if sampling_rate != TARGET_RESAMPLE: |
| filtered = np.reshape(filtered, (-1,1)) |
| filtered = librosa.resample(filtered, orig_sr=sampling_rate, target_sr=TARGET_RESAMPLE, axis=0, res_type='soxr_qq') |
| filtered = filtered.reshape(-1, TARGET_RESAMPLE*30) |
| |
| if cfg.save_type == 'npy': |
| |
| output_path = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn, sig) |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) |
| np.save(output_path, filtered.astype(np.float32)) |
| return |
| else: |
| return filtered |
|
|
|
|
| def process_eeg(args): |
| sig, chn_grp, chn = args |
| sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig) |
| data = np.load(sig_path) |
| sampling_rate = data.shape[1] // (30) |
| data = data.reshape(-1, 1) |
| |
| filtered = butter_bandpass_filter(data, EEG_LOWCUT, EEG_HIGHCUT, sampling_rate) |
| b_notch, a_notch = iirnotch(NOTCH_FREQ, 30, sampling_rate) |
| filtered = lfilter(b_notch, a_notch, filtered) |
| if sampling_rate != TARGET_RESAMPLE: |
| filtered = librosa.resample(filtered, orig_sr=sampling_rate, target_sr=TARGET_RESAMPLE, axis=0, res_type='soxr_qq') |
| filtered = filtered.reshape(-1, TARGET_RESAMPLE*30) |
| |
| if cfg.save_type == 'npy': |
| |
| output_path = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn, sig) |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) |
| np.save(output_path, filtered.astype(np.float32)) |
| return |
| else: |
| return filtered |
|
|
| def process_emg(args): |
| sig, chn_grp, chn = args |
| sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig) |
| data = np.load(sig_path) |
| sampling_rate = data.shape[1] // (30) |
| data = data.reshape(-1, 1) |
| |
| filtered = butter_bandpass_filter(data, EMG_LOWCUT, EMG_HIGHCUT, sampling_rate) |
| b_notch, a_notch = iirnotch(NOTCH_FREQ, 30, sampling_rate) |
| filtered = lfilter(b_notch, a_notch, filtered) |
| if sampling_rate != TARGET_RESAMPLE: |
| filtered = librosa.resample(filtered, orig_sr=sampling_rate, target_sr=TARGET_RESAMPLE, axis=0, res_type='soxr_qq') |
| filtered = filtered.reshape(-1, TARGET_RESAMPLE*30) |
| |
| if cfg.save_type == 'npy': |
| |
| output_path = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn, sig) |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) |
| np.save(output_path, filtered.astype(np.float32)) |
| return |
| else: |
| return filtered |
|
|
| def process_flow(args): |
| sig, chn_grp, chn = args |
| sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig) |
| data = np.load(sig_path) |
| sampling_rate = data.shape[1] // (30) |
| data = data.reshape(-1, 1) |
| |
| filtered = butter_bandpass_filter(data, FLOW_LOWCUT, FLOW_HIGHCUT, sampling_rate) |
| if sampling_rate != TARGET_RESAMPLE: |
| filtered = librosa.resample(filtered, orig_sr=sampling_rate, target_sr=TARGET_RESAMPLE, axis=0, res_type='soxr_qq') |
| |
| |
|
|
| filtered = filtered.reshape(-1, TARGET_RESAMPLE*30) |
| |
| if cfg.save_type == 'npy': |
| |
| output_path = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn, sig) |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) |
| np.save(output_path, filtered.astype(np.float32)) |
| return |
| else: |
| return filtered |
|
|
| def save_mmap(results, sig_list, chn_grp, chn, split='train'): |
| signame_list = [sig.split('.')[0] for sig in sig_list] |
| epoch_lengths = [sig.shape[0] for sig in results] |
| metadata = dict() |
| sig_info = dict() |
| start, end = 0, 0 |
| for signame, sig in zip(signame_list, results): |
| length = sig.shape[0] |
| end = start + length |
| info = dict() |
| info['pos'] = (start, end) |
| info['stats'] = {'sum': float(np.sum(sig)), 'sumsq': float(np.sum(sig ** 2))} |
| sig_info[signame] = info |
| start = end |
| |
|
|
| mmap_path = os.path.join(BASE_PATH, cfg.output_dir, split, chn_grp, f'{chn}.mmap') |
| signals_cat = np.concatenate(results, axis=0).astype(np.float32) |
| os.makedirs(os.path.dirname(mmap_path), exist_ok=True) |
| data = np.memmap(mmap_path, dtype=np.float32, mode='w+', shape=signals_cat.shape) |
| data[:] = signals_cat |
| data.flush() |
|
|
| metadata['sig_info'] = sig_info |
| metadata['data_shape'] = signals_cat.shape |
| info_path = os.path.join(BASE_PATH, cfg.output_dir, split, chn_grp, f'{chn}.pickle') |
| with open(info_path, 'wb') as f: |
| pickle.dump(metadata, f) |
|
|
|
|
| def preprocess(sig_list, chn_grp=None, chn=None, split='train'): |
| def _process(sig_list, chn_grp, chn, split): |
| print(f"Processing {chn_grp}/{chn}") |
| |
| |
| if chn_grp in ['EEG','EOG']: |
| |
| input_list = [(sig, chn_grp, chn) for sig in sig_list] |
| with multiprocessing.Pool(NUM_PROCESS) as pool: |
| results = list(tqdm(pool.imap(process_eeg, input_list), total=len(input_list))) |
| if cfg.save_type == 'mmap': |
| save_mmap(results, sig_list, chn_grp, chn, split) |
|
|
| elif chn_grp in ['EMG']: |
| |
| input_list = [(sig, chn_grp, chn) for sig in sig_list] |
| with multiprocessing.Pool(NUM_PROCESS) as pool: |
| results = list(tqdm(pool.imap(process_emg, input_list), total=len(input_list))) |
| if cfg.save_type == 'mmap': |
| save_mmap(results, sig_list, chn_grp, chn, split) |
|
|
| elif chn_grp in ['ECG']: |
|
|
| |
| input_list = [(sig, chn_grp, chn) for sig in sig_list] |
| with multiprocessing.Pool(NUM_PROCESS) as pool: |
| results = list(tqdm(pool.imap(process_ecg, input_list), total=len(input_list))) |
| if cfg.save_type == 'mmap': |
| save_mmap(results, sig_list, chn_grp, chn, split) |
|
|
| elif chn_grp in ['ABD', 'CHEST', 'AIRFLOW']: |
| |
| input_list = [(sig, chn_grp, chn) for sig in sig_list] |
| with multiprocessing.Pool(NUM_PROCESS) as pool: |
| results = list(tqdm(pool.imap(process_flow, input_list), total=len(input_list))) |
| if cfg.save_type == 'mmap': |
| save_mmap(results, sig_list, chn_grp, chn, split) |
|
|
|
|
| else: |
| results = [] |
| |
| for sig in tqdm(sig_list): |
| sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig) |
| data = np.load(sig_path) |
| sampling_rate = data.shape[1] // (30) |
| if sampling_rate != TARGET_RESAMPLE or chn == 'SpO2': |
| |
| data = data.reshape(-1, 1) |
| if chn == 'SpO2': |
| data = librosa.resample(data, orig_sr=sampling_rate, target_sr=1, axis=0, res_type='soxr_qq') |
| data = data / 1e8 |
| if np.max(data) > 1 or np.min(data) < 0: |
| print(f"Warning: SaO2 values out of range: {sig} with min {np.min(data)} and max {np.max(data)}") |
| data = data.reshape(-1, 1*30) |
| else: |
| data = librosa.resample(data, orig_sr=sampling_rate, target_sr=TARGET_RESAMPLE, axis=0, res_type='soxr_qq') |
| data = data.reshape(-1, TARGET_RESAMPLE*30) |
| |
| |
| if cfg.save_type == 'npy': |
| |
| output_path = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn, sig) |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) |
| np.save(output_path, data) |
| elif cfg.save_type == 'mmap': |
| results.append(data) |
| |
| save_mmap(results, sig_list, chn_grp, chn, split) |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| if chn_grp is not None and chn is not None: |
| _process(sig_list, chn_grp, chn, split) |
| return |
| else: |
| for chn_grp in os.listdir(os.path.join(BASE_PATH, cfg.data_dir)): |
| if not os.path.isdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp)): |
| continue |
| for chn in os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp)): |
| if not os.path.isdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)): |
| continue |
| _process(sig_list, chn_grp, chn, split) |
|
|
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument('--data_dir', default='#_MESA/signals_MAE', help='mros raw data location') |
| parser.add_argument('--output_dir', default='#_MESA/mmap_signals_MAE_processed', help='mros processed data location') |
| parser.add_argument('--ann_split', action='store_true', help='To split annotation files into train/val/test') |
| parser.add_argument('--save_type', default='mmap', help='Which data type to save. "npy" or "mmap"') |
| parser.add_argument('--chn_grp', default=None, help='If specified, only preprocess this channel group') |
| parser.add_argument('--chn', default=None, help='If specified, only preprocess this channel') |
| parser.add_argument('--mros2', action='store_true') |
|
|
| cfg = parser.parse_args() |
|
|
| if __name__ == '__main__': |
|
|
| if not cfg.mros2: |
| def get_mros_lists(mros_split_file): |
| ids = pd.read_csv(os.path.join("mesa_meta_data", mros_split_file), header=None)[0].astype(str).tolist() |
| return [x.strip() + ".npy" for x in ids] |
|
|
| list_files_train = get_mros_lists("1_train_patients.txt") |
| list_files_val = get_mros_lists("2_valid_patients.txt") |
| list_files_test = get_mros_lists("3_test_patients.txt") |
|
|
| sig_list = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, "EEG", "C4-M1"))) |
| |
| list_files_train = list(set(list_files_train).intersection(set(sig_list))) |
| list_files_val = list(set(list_files_val).intersection(set(sig_list))) |
| list_files_test = list(set(list_files_test).intersection(set(sig_list))) |
|
|
| print("Number of records in training set:", len(list_files_train)) |
| print("Number of records in validation set:", len(list_files_val)) |
| print("Number of records in test set:", len(list_files_test)) |
| else: |
| |
| sig_list = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, "EEG", "C4-M1"))) |
|
|
| |
| airflow_dir = os.path.join(BASE_PATH, cfg.data_dir, "AIRFLOW", "AIRFLOW") |
| airflow_files = set(os.listdir(airflow_dir)) |
| sig_list = sorted(list(set(sig_list).intersection(airflow_files))) |
|
|
| list_files_train = sig_list |
| list_files_val = [] |
| list_files_test = [] |
|
|
| print("Processing training set...") |
| preprocess(list_files_train, cfg.chn_grp, cfg.chn, split='train') |
| print("Processing evaluation set...") |
| preprocess(list_files_val, cfg.chn_grp, cfg.chn, split='val') |
| print("Processing test set...") |
| preprocess(list_files_test, cfg.chn_grp, cfg.chn, split='test') |
|
|
|
|
| |
| |
| if cfg.ann_split and not cfg.mros2: |
| ANN_PATH = os.path.join(os.path.dirname(os.path.join(BASE_PATH, cfg.data_dir)), 'annotations') |
| NEW_ANN_PATH = ANN_PATH.replace('annotations', 'annotations_split') |
| os.makedirs(NEW_ANN_PATH, exist_ok=True) |
| os.makedirs(os.path.join(NEW_ANN_PATH,'train'), exist_ok=True) |
| os.makedirs(os.path.join(NEW_ANN_PATH,'val'), exist_ok=True) |
| os.makedirs(os.path.join(NEW_ANN_PATH,'test'), exist_ok=True) |
|
|
| for fname in list_files_train: |
| ann_fname = fname.replace('.npy', '_ann.npy') |
| shutil.copy(os.path.join(ANN_PATH, ann_fname), os.path.join(NEW_ANN_PATH, 'train', ann_fname)) |
| for fname in list_files_val: |
| ann_fname = fname.replace('.npy', '_ann.npy') |
| shutil.copy(os.path.join(ANN_PATH, ann_fname), os.path.join(NEW_ANN_PATH, 'val', ann_fname)) |
| for fname in list_files_test: |
| ann_fname = fname.replace('.npy', '_ann.npy') |
| shutil.copy(os.path.join(ANN_PATH, ann_fname), os.path.join(NEW_ANN_PATH, 'test', ann_fname)) |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
|
|
| for split in ['train', 'val', 'test']: |
| |
| os.makedirs(os.path.join(BASE_PATH, cfg.output_dir, split, 'RESP'), exist_ok=True) |
| for chn_grp in ['ABD', 'AIRFLOW']: |
| chns = [f.split('.')[0] for f in os.listdir(os.path.join(BASE_PATH, cfg.output_dir, split, chn_grp))] |
| for chn in chns: |
| src_mmap = os.path.join(BASE_PATH, cfg.output_dir, split, chn_grp, f'{chn}.mmap') |
| src_pickle = os.path.join(BASE_PATH, cfg.output_dir, split, chn_grp, f'{chn}.pickle') |
| dst_mmap = os.path.join(BASE_PATH, cfg.output_dir, split, 'RESP', f'{chn}.mmap') |
| dst_pickle = os.path.join(BASE_PATH, cfg.output_dir, split, 'RESP', f'{chn}.pickle') |
| if not os.path.exists(dst_mmap): |
| os.symlink(src_mmap, dst_mmap) |
| if not os.path.exists(dst_pickle): |
| os.symlink(src_pickle, dst_pickle) |
| |
|
|
