data/process_mros1.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 = os.path.abspath('.')  # 현재 디렉토리 기준으로 저장
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_NAME = "Cannula Flow"

SELECT_CHNS = [
    "SaO2",
    ["ECG L-ECG R", ["ECG L", "ECG R"]],         
    ["L Chin-R Chin", ["L Chin", "R Chin"]], 
    ["EOG-Diff", ["LOC", "ROC"]],                   
    ["C3-A2", ["C3", "A2"]],    
    # ["C4-A1", ["C4", "A1"]],      
    "Abdominal", # ABD
    "Thoracic",  # CHEST
    AIRFLOW_NAME # AIRFLOW (없으면 skip)
]

# 저장 폴더 매핑(chn_grp / chn_name)
# (여기 key는 "내부에서 만든 신호 이름"임)
CHN_SAVE_DICT = {
    "SaO2": ("SaO2", "SaO2"),
    "ECG L-ECG R": ("ECG", "ECG"),
    "L Chin-R Chin": ("EMG", "EMG"),
    "EOG-Diff": ("EOG", "EOG-Diff"),
    "C3-A2": ("EEG", "C3-A2"),
    # "C4-A1": ("EEG", "C4-A1"),
    "Abdominal": ("ABD", "ABD"),
    "Thoracic": ("CHEST", "CHEST"),
    "Cannula Flow": ("AIRFLOW", "AIRFLOW"),
}




import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--data_dir', default="/home/yhna/data3/#_MrOS/polysomnography/edfs/visit1", help='MrOS EDF directory')
parser.add_argument('--annotation_dir', default="/home/yhna/data5/#_Annotation_XML/#_MrOS1_AHI_a0h3a", help='MrOS revised XML directory')
parser.add_argument('--output_dir', default="signals_MAE", help='Output directory under BASE_PATH (default: ./signals_MAE)')
parser.add_argument('--dataset_name', default='mros1', help='Folder name for saving annotations under BASE_PATH')
parser.add_argument('--num_workers', type=int, default=8)
cfg = parser.parse_args()

from tqdm.auto import tqdm

def process_signal(raw, select_idx, chns, fname):
    raw_sampling_rate = int(raw.info['sfreq'])
    raw_df = raw.to_data_frame()
    print(raw_df.columns)
    for chn_spec in chns:
    # chn_spec: str  또는  [a, [b,c]] 형태
        if isinstance(chn_spec, (list, tuple)):
            if len(chn_spec) != 2 or not isinstance(chn_spec[0], str):
                continue

            a = chn_spec[0]
            bc = chn_spec[1]

            if a in raw_df.columns:
                raw_ch_df = raw_df[a]
            else:
                if not (isinstance(bc, (list, tuple)) and len(bc) == 2):
                    continue
                b, c = bc[0], bc[1]
                if (b in raw_df.columns) and (c in raw_df.columns):
                    raw_df[a] = raw_df[b] - raw_df[c]
                    raw_ch_df = raw_df[a]
                else:
                    continue

            chn_name = a 

        else:
            chn_name = chn_spec
            if chn_name not in raw_df.columns:
                continue
            raw_ch_df = raw_df[chn_name]

        raw_ch_df = raw_ch_df.to_frame()
        raw_ch_df.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
        ep_len = EPOCH_SEC_SIZE * raw_sampling_rate
        if len(raw_ch) % ep_len != 0:
            print(f"Length of raw channel ({len(raw_ch)}) is not multiple of ({EPOCH_SEC_SIZE} * {raw_sampling_rate})")
            continue  

        # get (# of epochs)
        n_epochs = len(raw_ch) // ep_len 

        # split epochs
        x = raw_ch.reshape(n_epochs, ep_len).astype(np.float32)

        # select_idx 적용
        x = x[select_idx]
        x = x.reshape((-1, ep_len))

        # 저장 경로 결정: chn_name 기준으로 매핑
        chn_grp, chn_rename = CHN_SAVE_DICT[chn_name] 

        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
    ann_filename = os.path.basename(edf_fname).replace(".edf", "_ann.npy")
    labels = []


    # Read annotation and its header
    if ann_fname is None: 
        print(f"None of annfname: {edf_fname}")
        return -1

    # annotations-events-profusion -> annotations-events 변경으로 인한 
    t = ET.parse(ann_fname)
    r = t.getroot()
    faulty_File = 0

    use_r4 = False
    if len(list(r)) > 4:
        try:
            if len(r[4]) > 0 and (r[4][0].text is not None) and r[4][0].text.strip().isdigit():
                use_r4 = True
        except Exception:
            use_r4 = False

    if use_r4:
        for i in range(len(r[4])):
            lbl = int(r[4][i].text)
            if lbl == 4:
                labels.append(3)
            elif lbl == 5:
                labels.append(4)
            else:
                labels.append(lbl)
            if lbl > 5:
                faulty_File = 1

        if faulty_File == 1:
            print("Faulty labels found!!")
            return -1

        y = np.asarray(labels, dtype=np.int32)

    # --- (2) NSRR 타입(MrOS/SHHS nsrr): ScoredEvent(Stages|Stages)에서 구간을 epoch로 펼침 ---
    else:
        epoch_len = float(r.findtext(".//EpochLength") or EPOCH_SEC_SIZE)

        stage_events = []
        for ev in r.findall(".//ScoredEvent"):
            if (ev.findtext("EventType") or "").strip() != "Stages|Stages":
                continue
            concept = (ev.findtext("EventConcept") or "").strip()
            # concept: "Wake|0", "Stage 2 sleep|2", "REM sleep|5" 형태
            try:
                lbl = int(concept.split("|")[-1])
            except Exception:
                continue

            start = float(ev.findtext("Start") or 0.0)
            dur = float(ev.findtext("Duration") or 0.0)
            stage_events.append((start, dur, lbl))

        if len(stage_events) == 0:
            print(f"[SKIP] No stage events in XML: {ann_fname}")
            return -1

        stage_events.sort(key=lambda x: x[0])
        max_end = max(s + d for s, d, _ in stage_events)
        n_epochs = int(round(max_end / epoch_len))

        y = np.full(n_epochs, -1, dtype=np.int32)

        for start, dur, lbl in stage_events:
            # 원본 매핑 유지
            if lbl == 4:
                v = 3
            elif lbl == 5:
                v = 4
            else:
                v = lbl

            s_idx = int(round(start / epoch_len))
            n = int(round(dur / epoch_len))
            if n <= 0:
                continue
            y[s_idx:s_idx + n] = v

        # 빈 구간 채우기 (forward fill)
        last = 0
        for i in range(len(y)):
            if y[i] == -1:
                y[i] = last
            else:
                last = y[i]

    # Select on sleep periods
    w_edge_mins = 30
    nw_idx = np.where(y != 0)[0]                          # get index of first not WAKE
    start_idx = nw_idx[0] - (w_edge_mins * 2)             # start from 30min before first WAKE
    end_idx = nw_idx[-1] + (w_edge_mins * 2)              # end at 30min after first WAKE
    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]
    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, cfg.dataset_name, "annotations")
    os.makedirs(ann_dir, exist_ok=True)
    np.save(os.path.join(ann_dir, ann_filename), arr=y)


if __name__ == '__main__':
    START_IDX = 2300  # 이미 처리한 개수

    edf_dir = cfg.data_dir
    xml_dir = cfg.annotation_dir

    # 1) EDF 목록
    edf_fnames = sorted([
        os.path.join(edf_dir, f)
        for f in os.listdir(edf_dir)
        if f.lower().endswith(".edf")
    ])

    # 2) EDF -> XML 매칭: {base}.edf  <->  {base}-nsrr-revised.xml
    ann_fnames = []
    for edf_path in edf_fnames:
        base = os.path.splitext(os.path.basename(edf_path))[0]
        xml_path = os.path.join(xml_dir, f"{base}-nsrr-revised.xml")
        if not os.path.exists(xml_path):
            print(f"[SKIP] XML not found for {base}: {xml_path}")
            ann_fnames.append(None)
        else:
            ann_fnames.append(xml_path)

    # 3) input_list (XML 없는 건 스킵)
    input_list = [
        (ann_fnames[i], edf_fnames[i], i, False)
        for i in range(len(edf_fnames))
        if ann_fnames[i] is not None
    ]

    input_list = input_list[START_IDX:] 


    with multiprocessing.Pool(cfg.num_workers) as pool:
        list(tqdm(pool.imap(process_data, input_list), total=len(input_list)))