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