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