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__))), 'datasets') 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) # Apply bandpass filter 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 = librosa.resample(filtered, orig_sr=SAMPLING_RATE, target_sr=TARGET_RESAMPLE, axis=0, res_type='soxr_qq') 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': # Save processed data 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) # Apply bandpass filter 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': # Save processed data 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) # Apply bandpass filter 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': # Save processed data 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) # Apply bandpass filter 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') # median filter? filtered = filtered.reshape(-1, TARGET_RESAMPLE*30) if cfg.save_type == 'npy': # Save processed data 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 not (chn == "C4-A1" or chn == "AIRFLOW"): continue if chn_grp in ['EEG','EOG']: # sig_list = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn))) 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']: # sig_list = os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)) 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']: # sig_list = os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)) 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']: # sig_list = os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)) 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(sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)))): 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 == 'SaO2': # assuming 30s epochs data = data.reshape(-1, 1) if chn == 'SaO2': data = librosa.resample(data, orig_sr=sampling_rate, target_sr=1, axis=0, res_type='soxr_qq') print(np.min(data), np.max(data)) data = data / 1e8 # data range 0 to 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': # Save processed data 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='cfs/signals_MAE', help='SHHS raw data location') parser.add_argument('--output_dir', default='cfs/mmap_signals_MAE_processed', help='SHHS 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('--num_train', type=int, default=50, help='Number of training samples') cfg = parser.parse_args() if __name__ == '__main__': skip_val, skip_test = False, False if cfg.num_train <= 100: if os.path.exists(os.path.join(BASE_PATH, cfg.output_dir, 'val')): skip_val = True if os.path.exists(os.path.join(BASE_PATH, cfg.output_dir, 'test')): skip_test = True cfg.output_dir += f"_train{cfg.num_train}" def get_cfs_lists(cfs_split_file): ids = pd.read_csv(os.path.join("cfs_meta_data", cfs_split_file), header=None, names=['a']) ids = ids['a'].values.tolist() return [id + ".npy" for id in ids] list_files_train = get_cfs_lists("train.txt") list_files_val = get_cfs_lists("val.txt") list_files_test = get_cfs_lists("test.txt") sig_list = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, "EEG", "C3-M2"))) list_files_train = sorted(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))) if cfg.num_train <= 100: list_files_train = list_files_train[:cfg.num_train] 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)) print("Processing training set...") preprocess(list_files_train, cfg.chn_grp, cfg.chn, split='train') if skip_val: val_output_dir = cfg.output_dir val_output_dir = val_output_dir.replace(f"_train{cfg.num_train}","") os.symlink(os.path.join(BASE_PATH, val_output_dir, 'val'), os.path.join(BASE_PATH, cfg.output_dir, 'val')) else: print("Processing evaluation set...") preprocess(list_files_val, cfg.chn_grp, cfg.chn, split='val') if skip_test: test_output_dir = cfg.output_dir test_output_dir = test_output_dir.replace(f"_train{cfg.num_train}","") os.symlink(os.path.join(BASE_PATH, test_output_dir, 'test'), os.path.join(BASE_PATH, cfg.output_dir, 'test')) else: print("Processing test set...") preprocess(list_files_test, cfg.chn_grp, cfg.chn, split='test') # Reorganize annotations files # ANN_PATH = os.path.join(BASE_PATH, 'shhs1/annotations') if cfg.ann_split: 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 tqdm(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 tqdm(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 tqdm(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)) AH_PATH = os.path.join(os.path.dirname(os.path.join(BASE_PATH, cfg.data_dir)), 'ah_labels') NEW_AH_PATH = AH_PATH.replace('ah_labels', 'ah_labels_split') os.makedirs(NEW_AH_PATH, exist_ok=True) os.makedirs(os.path.join(NEW_AH_PATH,'train'), exist_ok=True) os.makedirs(os.path.join(NEW_AH_PATH,'val'), exist_ok=True) os.makedirs(os.path.join(NEW_AH_PATH,'test'), exist_ok=True) for fname in tqdm(list_files_train): ann_fname = fname.replace('.npy', '_ah.npy') shutil.copy(os.path.join(AH_PATH, ann_fname), os.path.join(NEW_AH_PATH, 'train', ann_fname)) for fname in tqdm(list_files_val): ann_fname = fname.replace('.npy', '_ah.npy') shutil.copy(os.path.join(AH_PATH, ann_fname), os.path.join(NEW_AH_PATH, 'val', ann_fname)) for fname in tqdm(list_files_test): ann_fname = fname.replace('.npy', '_ah.npy') shutil.copy(os.path.join(AH_PATH, ann_fname), os.path.join(NEW_AH_PATH, 'test', ann_fname)) for split in ['train', 'val', 'test']: # create link files to respiratory channels os.makedirs(os.path.join(BASE_PATH, cfg.output_dir, split, 'RESP'), exist_ok=True) for chn_grp in ['ABD', 'CHEST', '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)