Hugging Face's logo

gundo0102
/
SleepMaMi

English
Sleep
SleepFoundationModel
Model card Files
xet
 Community
SleepMaMi / data /physionet_preprocess.py
KeondoPark
Initial commit for public release
35f60b9
Raw History Blame Contribute Delete
12.3 kB
import os
import shutil
import numpy as np
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 librosa
import pickle
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
parser = argparse.ArgumentParser()
parser.add_argument('--data_dir', default='physionet_2018/signals_MAE', help='Physionet 2018 raw data location')
parser.add_argument('--output_dir', default='physionet_2018/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"')
cfg = parser.parse_args()
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)
# data = data.reshape(-1, 1)
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, 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)
# data = data.reshape(-1, 1)
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, 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)
# data = data.reshape(-1, 1)
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, 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)
# data = data.reshape(-1, 1)
data = data.reshape(-1,)
# Apply bandpass filter
filtered = butter_bandpass_filter(data, FLOW_LOWCUT, FLOW_HIGHCUT, SAMPLING_RATE)
# median filter?
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, 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):
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, 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, chn_grp, f'{chn}.pickle')
with open(info_path, 'wb') as f:
pickle.dump(metadata, f)
if __name__ == '__main__':
for chn_grp in os.listdir(os.path.join(BASE_PATH, cfg.data_dir)):
for chn in os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp)):
print(f"Processing {chn_grp}/{chn}")
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)
# sig_list = [sig.split('.')[0] for sig in sig_list]
# epoch_lengths = [sig.shape[0] for sig in results]
# sig_info = dict()
# start, end = 0, 0
# for sig, length in zip(sig_list, epoch_lengths):
# end = start + length
# sig_info[sig] = (start, end)
# start = end
# mmap_path = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, f'{chn}.mmap')
# signals_cat = np.concatenate(results, axis=0).astype(np.float32)
# data = np.memmap(mmap_path, dtype=np.float32, mode='w+', shape=signals_cat.shape)
# data[:] = signals_cat
# data.flush()
# info_path = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, f'{chn}.pickle')
# with open(info_path, 'wb') as f:
# pickle.dump(sig_info, f)
elif chn_grp in ['EMG']:
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_emg, input_list), total=len(input_list)))
if cfg.save_type == 'mmap':
save_mmap(results, sig_list, chn_grp, chn)
elif chn_grp in ['ECG']:
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_ecg, input_list), total=len(input_list)))
if cfg.save_type == 'mmap':
save_mmap(results, sig_list, chn_grp, chn)
elif chn_grp in ['ABD', 'CHEST', 'AIRFLOW']:
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_flow, input_list), total=len(input_list)))
if cfg.save_type == 'mmap':
save_mmap(results, sig_list, chn_grp, chn)
else:
results = []
for sig in tqdm(sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)))):
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':
data = librosa.resample(data, orig_sr=SAMPLING_RATE, target_sr=1, axis=0, res_type='soxr_qq')
data = data / 100 # data range 0 to 100
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, 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)
# output_path = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn, sig)
# shutil.copy(sig_path, output_path)
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)
sig_list = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)))
save_mmap(results, sig_list, chn_grp, chn)
# create link files to respiratory channels
os.makedirs(os.path.join(BASE_PATH, cfg.output_dir, '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, chn_grp))]
for chn in chns:
src_mmap = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, f'{chn}.mmap')
src_pickle = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, f'{chn}.pickle')
dst_mmap = os.path.join(BASE_PATH, cfg.output_dir, 'RESP', f'{chn}.mmap')
dst_pickle = os.path.join(BASE_PATH, cfg.output_dir, '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)