SleepMaMi / data /isruc_process_rawdata.py

KeondoPark

Initial commit for public release

35f60b9 5 months ago

7.13 kB

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)
            

	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)