SleepMaMi / data /kiss_preprocess.py

KeondoPark

Inference code added

a015518 4 months ago

17.8 kB

import os
import shutil
import numpy as np
import pandas as pd

from scipy.signal import iirnotch, lfilter
from data_utils import butter_bandpass_filter, correct_ecg_inversion

import matplotlib.pyplot as plt
from tqdm import tqdm
import shutil
import librosa
import pickle
import gc

BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'datasets')
SAMPLING_RATE = 200
TARGET_RESAMPLE = 100
NOTCH_FREQ = 60
EEG_LOWCUT = 0.3
EEG_HIGHCUT = 40
EMG_LOWCUT = 10
EMG_HIGHCUT = 1000
ECG_LOWCUT = 0.3
ECG_HIGHCUT = 70
FLOW_LOWCUT = 0
FLOW_HIGHCUT = 15

do_resample = (SAMPLING_RATE != TARGET_RESAMPLE)

NUM_PROCESS = 4

import argparse
import multiprocessing

def process_ecg(args):
    sig, chn_grp, chn, split = args
    sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig)
    data = np.load(sig_path)
    
    data = data.reshape(-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_epochs = []
    
    for epoch in filtered:
        inv_epoch, _ = correct_ecg_inversion(epoch, SAMPLING_RATE)
        filtered_epochs.append(inv_epoch)                
    
    filtered = np.stack(filtered_epochs, axis=0)

    if do_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, split, 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, split = args
    sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig)
    data = np.load(sig_path)    
    data = data.reshape(-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 do_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, split, 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, split = args
    sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig)
    data = np.load(sig_path)    
    data = data.reshape(-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 do_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, split, 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, split = args
    sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig)
    data = np.load(sig_path)    
    data = data.reshape(-1,)
    # Apply bandpass filter
    filtered = butter_bandpass_filter(data, FLOW_LOWCUT, FLOW_HIGHCUT, SAMPLING_RATE)    

    if do_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, split, 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)

    del data

# def get_kiss_lists(save_signals_path, snuh_split_file):
def get_kiss_lists(snuh_split_file):
    ids = pd.read_csv(os.path.join("kiss_meta_data", snuh_split_file), header=None, names=['a'])
    ids = ids['a'].values.tolist()    
    return [id + ".npy" for id in ids]
    # list_files = [os.path.join(save_signals_path, id + ".npy") for id in ids]
    # return list_files

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, split) 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)
            del results

        elif chn_grp in ['EMG']:            
            input_list = [(sig, chn_grp, chn, split) 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)
            del results

        elif chn_grp in ['ECG']:            
            input_list = [(sig, chn_grp, chn, split) 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)
            del results

        elif chn_grp in ['ABD', 'CHEST', 'AIRFLOW']:            
            input_list = [(sig, chn_grp, chn, split) 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)
            del results
        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)
                if do_resample or chn == 'SaO2':
                    data = data.reshape(-1, 1)
                    if chn == 'SaO2':
                        # When saving to "npy" (for inference testing), resample to 100 Hz. 
                        # When saving to "mmap" (for training), resample to 1 Hz.
                        if cfg.save_type == 'npy':
                            target_sr = 100
                        else:
                            target_sr = 1
                        data = librosa.resample(data, orig_sr=SAMPLING_RATE, target_sr=target_sr, axis=0, res_type='soxr_qq')
                        data = data / 100
                        if round(np.max(data)) > 1 or round(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, split, 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)
                    
            if cfg.save_type == 'mmap':
                save_mmap(results, sig_list, chn_grp, chn, split)
            del results
        gc.collect()
    
    
    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='kiss/signals_MAE', help='Physionet 2018 raw data location')
parser.add_argument('--output_dir', default='kiss/mmap_signals_MAE_processed', help='Physionet 2018 processed data location')
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('--kvss', action='store_true', help='KVSS preprocessing?')
parser.add_argument('--ann_split', action='store_true', help='To split annotation files into train/val/test')

cfg = parser.parse_args()


if __name__ == "__main__":
    
    if not cfg.kvss:
        list_files_train = sorted(get_kiss_lists("1_train_patients.txt"))
        list_files_val = sorted(get_kiss_lists("2_valid_patients.txt"))
        list_files_test = sorted(get_kiss_lists("3_test_patients.txt"))
        
        sig_list = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, "EEG", "C3-A2")))    
        
        for chn_grp in os.listdir(os.path.join(BASE_PATH, cfg.data_dir)):
            # Skip any files that are not directories
            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)):    
                # Skip any files that are not directories
                if not os.path.isdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)):
                    continue
                sig_list_chn = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)))
                print(f"Number of records in {chn_grp}/{chn}:", len(sig_list_chn))
                sig_list = list(set(sig_list).intersection(set(sig_list_chn)))
                
        print(f"Number of records that exist in all chns:", len(sig_list)) # Total 7698 -> 7579 records
        
        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)) # Total 6161 -> 6064 records
        print("Number of records in validation set:", len(list_files_val)) # Total 765 -> 748 records
        print("Number of records in test set:", len(list_files_test)) # Total 772 -> 767 records
    else:
        # For KVSS, use all data for preprocessing
        sig_list = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, "EEG", "C3-A2")))    
        
        for chn_grp in os.listdir(os.path.join(BASE_PATH, cfg.data_dir)):
            # Skip any files that are not directories
            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)):    
                # Skip any files that are not directories
                if not os.path.isdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)):
                    continue
                sig_list_chn = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)))
                print(f"Number of records in {chn_grp}/{chn}:", len(sig_list_chn))
                sig_list = list(set(sig_list).intersection(set(sig_list_chn)))
                
        print(f"Number of records that exist in all chns:", len(sig_list)) 
        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.kvss:
        # Reorganize annotations files
        ANN_PATH = os.path.join(BASE_PATH, 'kiss/annotations')
        NEW_ANN_PATH = os.path.join(BASE_PATH, 'kiss/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))

        AH_PATH = os.path.join(BASE_PATH, 'kiss/ah_labels')
        NEW_AH_PATH = os.path.join(BASE_PATH, 'kiss/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 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 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 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)

	import os
	import shutil
	import numpy as np
	import pandas as pd

	from scipy.signal import iirnotch, lfilter
	from data_utils import butter_bandpass_filter, correct_ecg_inversion

	import matplotlib.pyplot as plt
	from tqdm import tqdm
	import shutil
	import librosa
	import pickle
	import gc

	BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'datasets')
	SAMPLING_RATE = 200
	TARGET_RESAMPLE = 100
	NOTCH_FREQ = 60
	EEG_LOWCUT = 0.3
	EEG_HIGHCUT = 40
	EMG_LOWCUT = 10
	EMG_HIGHCUT = 1000
	ECG_LOWCUT = 0.3
	ECG_HIGHCUT = 70
	FLOW_LOWCUT = 0
	FLOW_HIGHCUT = 15

	do_resample = (SAMPLING_RATE != TARGET_RESAMPLE)

	NUM_PROCESS = 4

	import argparse
	import multiprocessing

	def process_ecg(args):
	sig, chn_grp, chn, split = args
	sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig)
	data = np.load(sig_path)

	data = data.reshape(-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_epochs = []

	for epoch in filtered:
	inv_epoch, _ = correct_ecg_inversion(epoch, SAMPLING_RATE)
	filtered_epochs.append(inv_epoch)

	filtered = np.stack(filtered_epochs, axis=0)

	if do_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, split, 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, split = args
	sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig)
	data = np.load(sig_path)
	data = data.reshape(-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 do_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, split, 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, split = args
	sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig)
	data = np.load(sig_path)
	data = data.reshape(-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 do_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, split, 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, split = args
	sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig)
	data = np.load(sig_path)
	data = data.reshape(-1,)
	# Apply bandpass filter
	filtered = butter_bandpass_filter(data, FLOW_LOWCUT, FLOW_HIGHCUT, SAMPLING_RATE)

	if do_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, split, 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)

	del data

	# def get_kiss_lists(save_signals_path, snuh_split_file):
	def get_kiss_lists(snuh_split_file):
	ids = pd.read_csv(os.path.join("kiss_meta_data", snuh_split_file), header=None, names=['a'])
	ids = ids['a'].values.tolist()
	return [id + ".npy" for id in ids]
	# list_files = [os.path.join(save_signals_path, id + ".npy") for id in ids]
	# return list_files

	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, split) 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)
	del results

	elif chn_grp in ['EMG']:
	input_list = [(sig, chn_grp, chn, split) 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)
	del results

	elif chn_grp in ['ECG']:
	input_list = [(sig, chn_grp, chn, split) 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)
	del results

	elif chn_grp in ['ABD', 'CHEST', 'AIRFLOW']:
	input_list = [(sig, chn_grp, chn, split) 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)
	del results
	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)
	if do_resample or chn == 'SaO2':
	data = data.reshape(-1, 1)
	if chn == 'SaO2':
	# When saving to "npy" (for inference testing), resample to 100 Hz.
	# When saving to "mmap" (for training), resample to 1 Hz.
	if cfg.save_type == 'npy':
	target_sr = 100
	else:
	target_sr = 1
	data = librosa.resample(data, orig_sr=SAMPLING_RATE, target_sr=target_sr, axis=0, res_type='soxr_qq')
	data = data / 100
	if round(np.max(data)) > 1 or round(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, split, 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)

	if cfg.save_type == 'mmap':
	save_mmap(results, sig_list, chn_grp, chn, split)
	del results
	gc.collect()


	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='kiss/signals_MAE', help='Physionet 2018 raw data location')
	parser.add_argument('--output_dir', default='kiss/mmap_signals_MAE_processed', help='Physionet 2018 processed data location')
	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('--kvss', action='store_true', help='KVSS preprocessing?')
	parser.add_argument('--ann_split', action='store_true', help='To split annotation files into train/val/test')

	cfg = parser.parse_args()


	if __name__ == "__main__":

	if not cfg.kvss:
	list_files_train = sorted(get_kiss_lists("1_train_patients.txt"))
	list_files_val = sorted(get_kiss_lists("2_valid_patients.txt"))
	list_files_test = sorted(get_kiss_lists("3_test_patients.txt"))

	sig_list = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, "EEG", "C3-A2")))

	for chn_grp in os.listdir(os.path.join(BASE_PATH, cfg.data_dir)):
	# Skip any files that are not directories
	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)):
	# Skip any files that are not directories
	if not os.path.isdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)):
	continue
	sig_list_chn = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)))
	print(f"Number of records in {chn_grp}/{chn}:", len(sig_list_chn))
	sig_list = list(set(sig_list).intersection(set(sig_list_chn)))

	print(f"Number of records that exist in all chns:", len(sig_list)) # Total 7698 -> 7579 records

	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)) # Total 6161 -> 6064 records
	print("Number of records in validation set:", len(list_files_val)) # Total 765 -> 748 records
	print("Number of records in test set:", len(list_files_test)) # Total 772 -> 767 records
	else:
	# For KVSS, use all data for preprocessing
	sig_list = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, "EEG", "C3-A2")))

	for chn_grp in os.listdir(os.path.join(BASE_PATH, cfg.data_dir)):
	# Skip any files that are not directories
	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)):
	# Skip any files that are not directories
	if not os.path.isdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)):
	continue
	sig_list_chn = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)))
	print(f"Number of records in {chn_grp}/{chn}:", len(sig_list_chn))
	sig_list = list(set(sig_list).intersection(set(sig_list_chn)))

	print(f"Number of records that exist in all chns:", len(sig_list))
	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.kvss:
	# Reorganize annotations files
	ANN_PATH = os.path.join(BASE_PATH, 'kiss/annotations')
	NEW_ANN_PATH = os.path.join(BASE_PATH, 'kiss/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))

	AH_PATH = os.path.join(BASE_PATH, 'kiss/ah_labels')
	NEW_AH_PATH = os.path.join(BASE_PATH, 'kiss/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 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 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 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)