data/process_mesa.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
AIRFLOW_CANDS = ["Flow"]  # strict cannula only

SELECT_CHNS = [
    ["SpO2", ["SpO2"]],
    ["ECG", ["ECG"]],
    ["EMG", ["EMG"]],

    ["EOG", [["EOG-L"], ["EOG-R"]]],
    ["Fz-Cz", ['EEG1']],
    ["Cz-Oz", ['EEG2']],
    ["C4-M1", ['EEG3']],

    ["Abdominal", ["Abdo"]],
    ["Thoracic",  ["Thor"]],

    ["Cannula Flow", AIRFLOW_CANDS],
]

CHN_SAVE_DICT = {
    "SpO2": ("SpO2", "SpO2"),
    "ECG": ("ECG", "ECG"),
    "EMG": ("EMG", "EMG"),
    "EOG": ("EOG", "EOG-Diff"),
    "Fz-Cz": ("EEG", "Fz-Cz"),
    "Cz-Oz": ("EEG", "Cz-Oz"),
    "C4-M1": ("EEG", "C4-M1"),
    "Abdominal": ("ABD", "ABD"),
    "Thoracic": ("CHEST", "CHEST"),
    "Cannula Flow": ("AIRFLOW", "AIRFLOW"),
}




import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--data_dir', default="/home/yhna/data3/#_MESA/polysomnography/edfs/", help='MESA EDF directory')
parser.add_argument('--annotation_dir', default="/home/yhna/data5/#_Annotation_XML/#_MESA_AHI_a0h3a", help='MESA 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='mesa', 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 _pick_first(raw_df, cands):
    if cands is None:
        return None, None
    if isinstance(cands, str):
        cands = [cands]
    for c in cands:
        if c in raw_df.columns:
            return raw_df[c], c
    return None, None


def process_signal(raw, select_idx, chns, fname):
    raw_sampling_rate = int(raw.info['sfreq'])
    raw_df = raw.to_data_frame()

    # ++++ 라벨(선택구간 포함) 기준으로 필요한 epoch 수 계산
    # ++++ select_idx는 원본 y의 인덱스이므로, 마지막 epoch까지 포함하려면 +1
    need_epochs = int(select_idx[-1] + 1)  # ++++
    ep_len = EPOCH_SEC_SIZE * raw_sampling_rate  # ++++
    need_len = need_epochs * ep_len  # ++++

    for chn_spec in chns:
        if isinstance(chn_spec, (list, tuple)):
            out_name, spec = chn_spec[0], chn_spec[1]

            # --- diff case: [[L cands], [R cands]] ---
            if (isinstance(spec, (list, tuple)) and len(spec) == 2
                  and isinstance(spec[0], (list, tuple)) and isinstance(spec[1], (list, tuple))):
                left, _ = _pick_first(raw_df, spec[0])
                right, _ = _pick_first(raw_df, spec[1])
                if left is None or right is None:
                    continue
                raw_df[out_name] = left - right
                raw_ch_df = raw_df[out_name]
                chn_name = out_name

            # --- alias case: [cand1, cand2, ...] ---
            else:
                raw_ch_df, _ = _pick_first(raw_df, spec)
                if raw_ch_df is None:
                    continue
                chn_name = out_name

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

        raw_ch = raw_ch_df.to_numpy().reshape(-1, 1)

        # ++++ 라벨이 존재하는 구간까지만 사용 (tail 1~29초 제거)
        # ++++ MESA에서는 label_epochs == floor_epochs 패턴이라 truncate가 정답
        if len(raw_ch) < need_len:  # ++++
            # ++++ 라벨/선택구간이 요구하는 길이보다 신호가 짧으면(이론상 MESA에서는 거의 없음)
            # ++++ 기존 정책대로 스킵
            print(f"[SKIP] Signal shorter than label-required length: {len(raw_ch)} < {need_len} ({fname}, {chn_name})")  # ++++
            continue  # ++++
        raw_ch = raw_ch[:need_len]  # ++++

        # ++++ 이제 무조건 30초 epoch 정수배가 됨
        n_epochs = need_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 = None  # 이미 처리한 개수

    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
    ]

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