SleepMaMi / data /cfs_process_rawdata.py
KeondoPark
Initial commit for public release
35f60b9
Raw History Blame Contribute Delete
11.1 kB
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')
# BASE_PATH = 'datasets'
EPOCH_SEC_SIZE = 30
# 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', 'SpO2','ECG1', 'EMG1', 'LOC', 'ROC', 'C3', 'C4', 'M1', 'M2', 'A1','A2', 'ABDO EFFORT', 'THOR EFFORT', 'AIRFLOW']
CHN_SAVE_DICT = {
'SaO2': ('SaO2','SaO2'),
'SpO2': ('SaO2','SaO2'),
'ECG1': ('ECG', 'ECG'),
'EMG1': ('EMG', 'EMG'),
'LOC': ('EOG', 'EOG-Diff'),
'ROC': ('EOG', 'EOG-Diff'),
'C3': ('EEG', 'C3-M2'),
'C4': ('EEG', 'C4-M1'),
'M1': ('EEG', 'C4-M1'),
'M2': ('EEG', 'C3-M2'),
'A1': ('EEG', 'A1-M2'),
'A2': ('EEG', 'A2-M1'),
'ABDO EFFORT': ('ABD', 'ABD'),
'THOR EFFORT': ('CHEST', 'CHEST'),
'AIRFLOW': ('AIRFLOW', 'AIRFLOW')
}
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--data_dir', default="CFS_raw/polysomnography/edfs", help='CFS raw data location')
parser.add_argument('--output_dir', default='cfs/signals_MAE', help='CFS processed data location')
parser.add_argument('--annotation_dir', default='CFS_raw/polysomnography/annotations-events-profusion', help='CFS annotations location')
parser.add_argument('--ah_dir', default='CFS_raw/polysomnography/annotations-events-nsrr', help='CFS apnea/hypopnea labels location')
parser.add_argument('--demographics_file', default='CFS_raw/datasets/cfs-visit5-harmonized-dataset-0.7.0.csv', help='CFS demographics file location')
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 == "LOC":
raw_df['EOG-Diff'] = raw_df['LOC'] - raw_df['ROC']
raw_ch_df = raw_df['EOG-Diff']
elif chn == "ROC":
continue
elif chn == 'C3':
if 'M2' in raw_df.columns:
raw_df['C3-M2'] = raw_df['C3'] - raw_df['M2']
else:
raw_df['C3-M2'] = raw_df['C3'] - raw_df['A2']
raw_ch_df = raw_df['C3-M2']
elif chn == 'C4':
if 'M1' in raw_df.columns:
raw_df['C4-M1'] = raw_df['C4'] - raw_df['M1']
else:
raw_df['C4-M1'] = raw_df['C4'] - raw_df['A1']
raw_ch_df = raw_df['C4-M1']
elif chn in ['M1', 'M2', 'A1', 'A2']:
continue
elif chn in ['SaO2', 'SpO2']:
if 'SaO2' in raw_df.columns:
raw_ch_df = raw_df['SaO2']
else:
raw_ch_df = raw_df['SpO2']
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))
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, ah_fname, edf_fname, file_id = 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, 'cfs', 'annotations')
os.makedirs(ann_dir, exist_ok=True)
np.save(os.path.join(ann_dir, ann_filename), arr=y)
tree = ET.parse(ah_fname)
root = tree.getroot()
events = root.findall('.//ScoredEvent')
total_seconds = 0
for event in events:
if event.find('EventConcept').text == 'Recording Start Time':
duration = float(event.find('Duration').text)
total_seconds = int(np.ceil(duration))
break
# 만약 Recording Start Time을 못 찾았을 경우 (예외 처리), 기존 방식대로 마지막 이벤트 시간으로 계산
if total_seconds == 0:
print("예외처리로 인한 Start time ~ End time 계산")
max_time = 0
for event in events:
try:
start = float(event.find('Start').text)
dur = float(event.find('Duration').text)
if start + dur > max_time:
max_time = start + dur
except:
continue
total_seconds = int(np.ceil(max_time))
# 2. 라벨 어레이 생성
ah_labels = np.zeros(total_seconds, dtype=int)
# 3. 호흡 이벤트 파싱 및 라벨링
for event in events:
try:
event_type = event.find('EventType').text
if event_type != 'Respiratory|Respiratory': continue # 호흡 이벤트가 아니면 스킵
concept = event.find('EventConcept').text
start = float(event.find('Start').text)
duration = float(event.find('Duration').text)
start_idx = int(round(start))
end_idx = start_idx + int(round(duration))
# 인덱스 범위 보정
if end_idx > total_seconds:
end_idx = total_seconds
# 라벨링 (Hypopnea=1, Apnea=2)
if 'Hypopnea' in concept:
mask = ah_labels[start_idx:end_idx] != 2
ah_labels[start_idx:end_idx][mask] = 1
elif any(x in concept for x in ['Obstructive', 'Central', 'Mixed']):
ah_labels[start_idx:end_idx] = 2
except AttributeError:
continue
ah_labels = np.reshape(ah_labels, (-1, 30))
ah_labels = ah_labels[select_idx]
assert len(ah_labels) == len(y), f"Length of ah_labels ({len(ah_labels)}) does not match length of sleep stage labels ({len(y)})"
ah_labels = np.reshape(ah_labels, (-1,))
ah_dir = os.path.join(BASE_PATH, 'cfs', 'ah_labels')
ah_filename = os.path.basename(edf_fname).replace(".edf", "_ah.npy")
os.makedirs(ah_dir, exist_ok=True)
np.save(os.path.join(ah_dir, ah_filename), arr=ah_labels)
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))
demo = pd.read_csv(os.path.join(BASE_PATH, cfg.demographics_file))
def get_nsrrid(filename):
nsrrid = filename.split('-')[-1].replace('.edf', '')
return int(nsrrid)
age_filtered_df = demo[demo['nsrr_age'] > 13]
file_list = [filename for filename in file_list if get_nsrrid(filename) in age_filtered_df['nsrrid'].values]
print(f"Number of files after filtering by age > 13: {len(file_list)}")
filenames_edf = [filename for filename in file_list if filename.endswith('.edf')]
filenames_ann = [filename.replace('.edf', '-profusion.xml') for filename in filenames_edf]
filenames_ah = [filename.replace('.edf', '-nsrr.xml') for filename in filenames_edf]
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]
ah_fnames = [os.path.join(BASE_PATH, cfg.ah_dir, filename_ah) for filename_ah in filenames_ah]
edf_fnames.sort() # sort edf file names list
ann_fnames.sort() # sort ann file names list
ah_fnames.sort() # sort ah file names list
input_list = zip(ann_fnames, ah_fnames, edf_fnames, range(len(ann_fnames)))
# 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)))