File size: 7,131 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 | import scipy.io
import os
from tqdm import tqdm
import numpy as np
import mne
import shutil
CHN_SAVE_DICT = {
'C3-A2': ('EEG', 'C3-A2'),
'C4-A1': ('EEG', 'C4-A1'),
'F3-A2': ('EEG', 'F3-A2'),
'F4-A1': ('EEG', 'F4-A1'),
'O1-A2': ('EEG', 'O1-A2'),
'O2-A1': ('EEG', 'O2-A1'),
'C3-M2': ('EEG', 'C3-A2'),
'C4-M1': ('EEG', 'C4-A1'),
'F3-M2': ('EEG', 'F3-A2'),
'F4-M1': ('EEG', 'F4-A1'),
'O1-M2': ('EEG', 'O1-A2'),
'O2-M1': ('EEG', 'O2-A1'),
'LOC-A2': ('EOG', 'Left-A2'),
'ROC-A1': ('EOG', 'Right-A1'),
'E1-M2': ('EOG', 'Left-A2'),
'E2-M1': ('EOG', 'Right-A1'),
'X1': ('EMG', 'EMG'),
'X2': ('ECG', 'ECG'),
'X6': ('AIRFLOW', 'AIRFLOW'),
'X7': ('CHEST', 'CHEST'),
'X8': ('ABD', 'ABD'),
'24': ('EMG', 'EMG'),
'25': ('ECG', 'ECG'),
'29': ('AIRFLOW', 'AIRFLOW'),
'30': ('CHEST', 'CHEST'),
'31': ('ABD', 'ABD'),
'SpO2': ('SaO2', 'SaO2'),
'SaO2': ('SaO2', 'SaO2')
}
SELECT_CHNS1 = ['F3-A2', 'F4-A1', 'C3-A2', 'C4-A1', 'O1-A2', 'O2-A1', 'LOC-A2', 'ROC-A1', 'X1', 'X2', 'X6', 'X7', 'X8', 'SaO2']
SELECT_CHNS2 = ['F3-M2', 'F4-M1', 'C3-M2', 'C4-M1', 'O1-M2', 'O2-M1', 'E1-M2', 'E2-M1', '24', '25', '29', '30', '31', 'SpO2']
# def match_channel_group(chn):
# """
# Find the group of the channel.
# """
# chn_group = {
# 'EEG': ['F3-A2', 'F4-A1', 'C3-A2', 'C4-A1', 'O1-A2', 'O2-A1'],
# 'EOG': ['LOC-A2', 'ROC-A1'],
# 'EMG': ['X1'],
# 'ABD': ['X8'],
# 'CHEST': ['X7'],
# 'AIRFLOW': ['X6'],
# 'SaO2': ['SaO2'],
# 'ECG': ['X2']
# }
# for k, v in chn_group.items():
# if chn in v:
# return k
# return None
BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), 'datasets')
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--data_dir', default='ISRUC', help='ISRUC raw data location')
parser.add_argument('--output_dir', default='ISRUC', help='ISRUC processed data location')
parser.add_argument('--annotation_dir', default='ISRUC', help='ISRUC annotations location')
cfg = parser.parse_args()
def read_ann(ann_path):
# Annotation
with open(ann_path, 'r') as f:
ann = f.readlines()
ann = np.array([line.strip() for line in ann if line.strip()]).astype(np.int64)
ann = np.where(ann==5, 4, ann)
return ann
# Example usage
raw_sampling_rate = 200
EPOCH_SEC_SIZE = 30
for ds_id in [1,2,3]:
if ds_id == 1:
num_subject = 100
elif ds_id == 2:
num_subject = 8
else:
num_subject = 10
num_session = 2 if ds_id == 2 else 1
for i in tqdm(range(1, num_subject+1)):
for j in range(1, num_session+1):
data_dir_full = os.path.join(BASE_PATH, cfg.data_dir, f'SG{ds_id}', f'{i}',f'{i}')
if ds_id == 2:
file_path = os.path.join(data_dir_full, f'{j}', f'{j}.edf')
ann_path = os.path.join(data_dir_full, f'{j}', f'{j}_1.txt')
if not os.path.exists(file_path):
rarpath = os.path.join(data_dir_full, f'{j}', f'{j}.rec')
shutil.copy(rarpath, file_path)
else:
file_path = os.path.join(data_dir_full,f'{i}.edf')
ann_path = os.path.join(data_dir_full,f'{i}_1.txt')
if not os.path.exists(file_path):
rarpath = os.path.join(data_dir_full, f'{i}.rec')
shutil.copy(rarpath, file_path)
ann = read_ann(ann_path)
raw = mne.io.read_raw_edf(file_path, preload=True)
raw_df = raw.to_data_frame()
if not len(raw_df) == len(ann) * EPOCH_SEC_SIZE * raw_sampling_rate:
print(f"Warning: {i} signal length and annotation length mismatched")
# Try to use the second annotator's scoring
if ds_id == 2:
ann_path = os.path.join(data_dir_full, f'{j}', f'{j}_2.txt')
else:
ann_path = os.path.join(data_dir_full,f'{i}_2.txt')
ann = read_ann(ann_path)
if not len(raw_df) == len(ann) * EPOCH_SEC_SIZE * raw_sampling_rate:
print(f"Error both annotations mismatch signal length, skip this record. " , "Dataset id:", ds_id, "Subject id:", i)
continue
ann_dir = os.path.join(BASE_PATH, f'{cfg.output_dir}_SG{ds_id}', 'annotations')
os.makedirs(ann_dir, exist_ok=True)
fname = f'{i:03}_{j}' if ds_id == 2 else f'{i:03}'
ann_name = f'{i:03}_{j}_ann' if ds_id == 2 else f'{i:03}_ann'
np.save(os.path.join(ann_dir, ann_name), ann)
for chn in SELECT_CHNS1:
# chn_grp = match_channel_group(chn)
if chn in raw_df.columns:
raw_ch_df = raw_df[chn]
chn_grp, chn_rename = CHN_SAVE_DICT[chn]
else:
idx = SELECT_CHNS1.index(chn)
chn = SELECT_CHNS2[idx]
if chn in raw_df.columns:
raw_ch_df = raw_df[chn]
chn_grp, chn_rename = CHN_SAVE_DICT[chn]
else:
print(f"Channel {chn} not found in {file_path}, skip this channel.")
print(raw_df.columns)
continue
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_df) % (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})")
# 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))
os.makedirs(os.path.join(BASE_PATH, f'{cfg.output_dir}_SG{ds_id}', 'signals_MAE', chn_grp, chn_rename), exist_ok=True)
np.save(os.path.join(BASE_PATH, f'{cfg.output_dir}_SG{ds_id}', 'signals_MAE', chn_grp, chn_rename, f"{fname}.npy"), x)
|