SleepMaMi / data /shhs_process_rawdata.py

KeondoPark

Initial commit for public release

35f60b9 5 months ago

8.12 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')
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)))

	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)))