data/cfs_process_rawdata.py

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