File size: 8,123 Bytes
35f60b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 | 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')
EPOCH_SEC_SIZE = 30
# ALL_CHANNELS = ['SaO2', 'H.R.', 'EEG(sec)', 'ECG', 'EMG', 'EOG(L)', 'EOG(R)', 'EEG', 'THOR RES', 'ABDO RES', 'POSITION', 'LIGHT', 'NEW AIR', 'OX stat']
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', 'ECG', 'EMG', 'EOG(L)', 'EOG(R)', 'EEG', 'ABDO RES'] + AIRFLOW_NAMES + EEG2_NAMES
CHN_SAVE_DICT = {
'SaO2': ('SaO2','SaO2'),
# 'EEG(sec)': ('EEG', 'C4-A1'),
# 'EEG2': ('EEG', 'C4-A1'),
# 'EEG 2': ('EEG', 'C4-A1'),
# 'EEG sec': ('EEG', 'C4-A1'),
'ECG': ('ECG', 'ECG'),
'EMG': ('EMG', 'EMG'),
'EOG(L)': ('EOG', 'EOG-Diff'),
'EOG(R)': ('EOG', 'EOG-Diff'),
'EEG': ('EEG', 'C3-A2'),
'ABDO RES': ('ABD', 'ABD'),
'NEW AIR': ('AIRFLOW', 'AIRFLOW'),
'AIRFLOW': ('AIRFLOW', 'AIRFLOW')
}
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--data_dir', default="#_SHHS/polysomnography/edfs/shhs1", help='SHHS1 raw data location')
parser.add_argument('--output_dir', default='shhs1/signals_MAE', help='SHHS1 processed data location')
parser.add_argument('--annotation_dir', default='#_SHHS/polysomnography/annotations-events-profusion/shhs1', help='SHHS1 annotations location')
parser.add_argument('--shhs2', action='store_true', help='Use SHHS2 dataset')
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 == "EOG(L)":
raw_df['EOG-Diff'] = raw_df['EOG(L)'] - raw_df['EOG(R)']
raw_ch_df = raw_df['EOG-Diff']
elif chn == "EOG(R)":
continue
elif chn in AIRFLOW_NAMES:
if chn in raw_df.columns:
raw_ch_df = raw_df[chn]
else:
continue
elif chn in EEG2_NAMES:
if chn in raw_df.columns:
raw_ch_df = raw_df[chn]
else:
continue
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))
if chn in AIRFLOW_NAMES:
chn_grp, chn_rename = 'AIRFLOW', 'AIRFLOW'
elif chn in EEG2_NAMES:
chn_grp, chn_rename = 'EEG', 'C4-A1'
else:
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, edf_fname, file_id, shhs2 = 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, 'shhs1' if not shhs2 else 'shhs2', "annotations")
os.makedirs(ann_dir, exist_ok=True)
np.save(os.path.join(ann_dir, ann_filename), arr=y)
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))
file_list = [f for f in file_list if f.startswith('tr')]
filenames_edf = pd.read_csv("shhs_meta_data/filenames_edf.txt", header=None, names=['a'])
filenames_edf = filenames_edf['a'].values.tolist()
filenames_ann = pd.read_csv("shhs_meta_data/filenames_ann.txt", header=None, names=['a'])
filenames_ann = filenames_ann['a'].values.tolist()
if cfg.shhs2:
edf_fnames = [os.path.join(BASE_PATH, cfg.data_dir, f) for f in os.listdir(os.path.join(BASE_PATH, cfg.data_dir))]
ann_fnames = [os.path.join(BASE_PATH, cfg.annotation_dir, f) for f in os.listdir(os.path.join(BASE_PATH, cfg.annotation_dir))]
else:
# make actual file names for selected edf files and make a list
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]
edf_fnames.sort() # sort edf file names list
ann_fnames.sort() # sort ann file names list
# sorted edf/ann file names list to nparray
# edf_fnames = np.asarray(edf_fnames)
# ann_fnames = np.asarray(ann_fnames)
input_list = zip(ann_fnames, edf_fnames, range(len(ann_fnames)), len(ann_fnames) * [cfg.shhs2])
# 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)))
|