Add code: rest_prep_1020.py

code/rest_prep_1020.py

@@ -0,0 +1,251 @@
+import os, re, math, json, argparse, warnings
+from pathlib import Path
+from typing import List, Tuple, Dict, Optional
+import numpy as np, h5py, pyedflib
+
+# ---------- 10-20 electrode 2D coords (x,y) đã chuẩn hoá tương đối ----------
+E1020 = {
+ "FP1":(-0.5,  1.0), "FP2":( 0.5,  1.0),
+ "F7": (-1.0,  0.6), "F3":(-0.3,  0.6), "FZ":(0.0,0.8), "F4":( 0.3,  0.6), "F8":( 1.0,  0.6),
+ "T3":(-1.2,  0.2), "T7":(-1.2,  0.2), "C3":(-0.4,  0.2), "CZ":(0.0,0.3), "C4":( 0.4,  0.2), "T4":( 1.2,  0.2), "T8":( 1.2,  0.2),
+ "T5":(-1.1, -0.2), "P7":(-1.1, -0.2), "P3":(-0.3, -0.2), "PZ":(0.0,-0.1), "P4":( 0.3, -0.2), "T6":( 1.1, -0.2), "P8":( 1.1, -0.2),
+ "O1":(-0.5, -0.8), "O2":( 0.5, -0.8)
+}
+ALIASES = {"T3":"T7","T4":"T8","T5":"P7","T6":"P8"}  # đồng nhất tên
+
+_BLOCKLIST = re.compile(r'(?:ECG|VNS|-$)', re.IGNORECASE)
+
+def _norm_e(name:str)->Optional[str]:
+    name=name.upper().strip()
+    name=ALIASES.get(name,name)
+    return name if name in E1020 else None
+
+def pair_midpoint(ch_label:str)->Optional[Tuple[float,float]]:
+    """
+    CHB-MIT thường dùng kênh lưỡng cực 'A-B'. Lấy toạ độ là trung điểm của hai điện cực A,B.
+    """
+    if '-' not in ch_label: return None
+    a,b = [x.strip().upper() for x in ch_label.split('-',1)]
+    a=_norm_e(a); b=_norm_e(b)
+    if (a is None) or (b is None): return None
+    ax,ay=E1020[a]; bx,by=E1020[b]
+    return ((ax+bx)/2.0, (ay+by)/2.0)
+
+def build_1020_graph(channels:List[str], k:int=8, sigma:float=0.5, radius:float=1.6):
+    """
+    Đồ thị khoảng cách tĩnh theo 10–20:
+      w_ij = exp(-||vi-vj||^2 / sigma^2), nếu ||vi-vj|| <= radius, ngược lại 0.
+      Sau đó với mỗi nút, giữ lại k láng giềng w lớn nhất (undirected).
+    """
+    coords=[]
+    keep=[]
+    for ch in channels:
+        if _BLOCKLIST.search(ch): continue
+        m=pair_midpoint(ch)
+        if m is not None:
+            coords.append(m); keep.append(ch)
+    C=len(keep)
+    if C<2:
+        raise RuntimeError("Không đủ kênh ánh xạ được sang 10-20")
+    coords=np.array(coords, dtype=np.float32)
+    D=np.sqrt(((coords[None,:,:]-coords[:,None,:])**2).sum(axis=-1))  # (C,C)
+    W=np.exp(-(D**2)/(sigma**2))
+    W[D>radius]=0.0
+    np.fill_diagonal(W, 0.0)
+
+    edges=set()
+    for i in range(C):
+        idx=np.argsort(-W[i])[:max(1,min(k,C-1))]
+        for j in idx:
+            a,b=(i,j) if i<j else (j,i)
+            if W[a,b]>0: edges.add((a,b))
+    edges=sorted(list(edges))
+    ei=np.array(edges, dtype=np.int64).T
+    ei=np.hstack([ei, ei[::-1,:]]) if ei.size else ei
+    ew=np.array([W[i,j] for (i,j) in edges], dtype=np.float32)
+    ew=np.concatenate([ew, ew]) if ew.size else ew
+    return keep, ei, ew
+
+# ---------- EDF utils ----------
+def list_edf_files(patient_dir: Path):
+    return sorted([p for p in patient_dir.glob("*.edf") if p.is_file()])
+
+def edf_channel_labels(edf_path: Path):
+    f=pyedflib.EdfReader(str(edf_path))
+    labels=[f.getLabel(i).strip() for i in range(f.signals_in_file)]
+    f._close(); del f
+    return [ch for ch in labels if not _BLOCKLIST.search(ch)]
+
+def intersection_channels(edf_paths):
+    common=None
+    for p in edf_paths:
+        chans=set(edf_channel_labels(p))
+        if not chans: continue
+        common = chans if common is None else (common & chans)
+    return sorted(list(common)) if common else []
+
+def read_edf_signals(edf_path: Path, keep_channels):
+    f=pyedflib.EdfReader(str(edf_path))
+    labels=[f.getLabel(i).strip() for i in range(f.signals_in_file)]
+    fs_all=[int(round(f.getSampleFrequency(i))) for i in range(f.signals_in_file)]
+    fs=int(round(np.median(fs_all))) if fs_all else 256
+    idxs=[]; out_labels=[]
+    for ch in keep_channels:
+        try:
+            i=labels.index(ch)
+            if _BLOCKLIST.search(ch): continue
+            idxs.append(i); out_labels.append(ch)
+        except ValueError:
+            f._close(); del f
+            raise RuntimeError(f"Channel {ch} not found in {edf_path.name}")
+    sigs=np.vstack([f.readSignal(i) for i in idxs]).astype(np.float32)  # (C,N)
+    f._close(); del f
+    return sigs, fs, out_labels
+
+# ---------- nhãn co giật ----------
+def parse_summary(summary_path: Path) -> Dict[str, List[Tuple[float,float]]]:
+    mapping={}
+    if not (summary_path and summary_path.exists()):
+        return mapping
+    curr=None; buf=[]
+    with summary_path.open("r", errors="ignore") as f:
+        for line in f:
+            line=line.strip()
+            mfile=re.search(r'File Name:\s*(\S+\.edf)', line, re.IGNORECASE)
+            if mfile:
+                if curr and buf:
+                    pairs=[(buf[i], buf[i+1]) for i in range(0,len(buf)-1,2)]
+                    mapping.setdefault(curr, []).extend(pairs)
+                curr=mfile.group(1); buf=[]; continue
+            if re.search(r'Seizure (Start|End) Time', line, re.IGNORECASE):
+                nums=[float(x) for x in re.findall(r'[\d.]+', line)]
+                if nums: buf.extend(nums)
+    if curr and buf:
+        pairs=[(buf[i], buf[i+1]) for i in range(0,len(buf)-1,2)]
+        mapping.setdefault(curr, []).extend(pairs)
+    return mapping
+
+def parse_seizures_file(seiz_file: Path) -> List[Tuple[float,float]]:
+    intervals=[]
+    if not seiz_file.exists(): return intervals
+    with seiz_file.open("r", errors="ignore") as f:
+        for line in f:
+            nums=[float(x) for x in re.findall(r'[-+]?\d*\.?\d+', line)]
+            if len(nums)>=2: intervals.append((nums[0], nums[1]))
+    return intervals
+
+def slice_starts(N, fs, clip_sec, hop_sec):
+    T=int(fs*clip_sec); hop=int(fs*hop_sec)
+    if N<T: return np.zeros((0,), dtype=np.int64)
+    return np.arange(0, N-T+1, hop, dtype=np.int64)
+
+def zscore_perclip(clip):
+    mu=clip.mean(axis=1, keepdims=True); sd=clip.std(axis=1, keepdims=True)+1e-8
+    return (clip-mu)/sd
+
+def label_for_window(start, T, fs, intervals, min_overlap_sec):
+    a,b=start, start+T
+    thr=int(round(min_overlap_sec*fs)); overlap=0
+    for (u,v) in intervals:
+        u_s=int(round(u*fs)); v_s=int(round(v*fs))
+        overlap += max(0, min(b, v_s) - max(a, u_s))
+        if overlap>=thr: return 1
+    return 0
+
+class H5Appender:
+    def __init__(self, out_path: Path, C: int, T: int, fs: int, channels, edge_index, edge_weight, gzip_level=4):
+        self.f=h5py.File(str(out_path),"w")
+        self.ds_clips=self.f.create_dataset("clips", shape=(0,C,T,1), maxshape=(None,C,T,1),
+            dtype="float32", chunks=(16,C,T,1), compression="gzip", compression_opts=gzip_level)
+        self.ds_labels=self.f.create_dataset("labels", shape=(0,), maxshape=(None,), dtype="i8",
+            chunks=True, compression="gzip", compression_opts=gzip_level)
+        self.ds_fileids=self.f.create_dataset("file_ids", shape=(0,), maxshape=(None,), dtype="i4",
+            chunks=True, compression="gzip", compression_opts=gzip_level)
+        self.f.attrs["fs"]=fs; self.f.attrs["T"]=T; self.f.attrs["patient"]=out_path.stem
+        self.f.create_dataset("channels", data=np.array([c.encode() for c in channels]))
+        self.f.create_dataset("edge_index", data=edge_index.astype(np.int64))
+        self.f.create_dataset("edge_weight", data=edge_weight.astype(np.float32))
+        self.n=0
+    def append(self, clips_CT, labels, file_id:int):
+        if clips_CT.size==0: return
+        M,C,T=clips_CT.shape
+        clips=clips_CT[...,None].astype(np.float32)
+        self.ds_clips.resize(self.n+M, axis=0)
+        self.ds_labels.resize(self.n+M, axis=0)
+        self.ds_fileids.resize(self.n+M, axis=0)
+        self.ds_clips[self.n:self.n+M]=clips
+        self.ds_labels[self.n:self.n+M]=labels.astype(np.int64)
+        self.ds_fileids[self.n:self.n+M]=np.full((M,), file_id, dtype=np.int32)
+        self.n+=M
+    def close(self): self.f.close()
+
+def process_patient_1020(root: Path, patient: str, out_path: Path,
+                         clip_sec: float=4.0, hop_sec: float=2.0,
+                         fs_target: int=256, min_overlap_sec: float=0.25,
+                         graph_k: int=8, sigma: float=0.5, radius: float=1.6):
+    pat_dir=root/patient
+    assert pat_dir.exists(), f"Not found: {pat_dir}"
+    edfs=list_edf_files(pat_dir); assert edfs, f"No EDF in {pat_dir}"
+
+    # kênh giao nhau trong bệnh nhân
+    keep_channels=intersection_channels(edfs)
+    # chỉ giữ kênh map được sang 10-20 (bỏ kênh lạ)
+    keep_channels=[ch for ch in keep_channels if pair_midpoint(ch) is not None]
+    if len(keep_channels)<8:
+        warnings.warn(f"[{patient}] chỉ còn {len(keep_channels)} kênh sau khi map 10-20")
+
+    # đồ thị 10-20 tĩnh
+    keep_channels, ei, ew = build_1020_graph(keep_channels, k=graph_k, sigma=sigma, radius=radius)
+
+    summ = next(iter(list(pat_dir.glob("*summary*.txt"))), None)
+    summ_map=parse_summary(summ) if summ else {}
+    seiz_map={}
+    for p in edfs:
+        iv=parse_seizures_file(p.with_suffix(p.suffix+".seizures"))
+        if iv: seiz_map[p.name]=iv
+    def intervals_for(name): return seiz_map.get(name, summ_map.get(name, []))
+
+    app=None; file_id=0
+    total_pos=0; total=0
+
+    for edf in edfs:
+        sigs, fs, chans=read_edf_signals(edf, keep_channels)
+        # resample
+        if fs!=fs_target:
+            ratio=fs_target/fs
+            N_new=int(round(sigs.shape[1]*ratio))
+            t_old=np.linspace(0, sigs.shape[1]-1, sigs.shape[1], dtype=np.float32)
+            t_new=np.linspace(0, sigs.shape[1]-1, N_new, dtype=np.float32)
+            sigs=np.stack([np.interp(t_new, t_old, ch) for ch in sigs], axis=0).astype(np.float32)
+
+        C,N=sigs.shape; T=int(fs_target*clip_sec)
+        starts=slice_starts(N, fs_target, clip_sec, hop_sec)
+        if app is None:
+            app=H5Appender(out_path, C=C, T=T, fs=fs_target, channels=chans, edge_index=ei, edge_weight=ew)
+        ivals=intervals_for(edf.name)
+        labels=np.array([label_for_window(int(s), T, fs_target, ivals, min_overlap_sec) for s in starts], dtype=np.int64)
+        M=len(starts); clips=np.empty((M,C,T), dtype=np.float32)
+        for i,s in enumerate(starts): clips[i]=zscore_perclip(sigs[:, s:s+T])
+        app.append(clips, labels, file_id=file_id)
+        total += M; total_pos += int(labels.sum()); file_id+=1
+
+    app.close()
+    print(f"[Done] {patient} -> {out_path} | clips={total} (pos={total_pos}, neg={total-total_pos}), C={len(keep_channels)}, T={int(fs_target*clip_sec)}")
+
+if __name__=="__main__":
+    import argparse
+    ap=argparse.ArgumentParser()
+    ap.add_argument("--root", required=True)
+    ap.add_argument("--patient", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--clip-sec", type=float, default=4.0)
+    ap.add_argument("--hop-sec", type=float, default=2.0)
+    ap.add_argument("--fs", type=int, default=256)
+    ap.add_argument("--min-overlap", type=float, default=0.25)
+    ap.add_argument("--k", type=int, default=8)
+    ap.add_argument("--sigma", type=float, default=0.5)
+    ap.add_argument("--radius", type=float, default=1.6)
+    args=ap.parse_args()
+    process_patient_1020(Path(args.root), args.patient, Path(args.out),
+        clip_sec=args.clip_sec, hop_sec=args.hop_sec, fs_target=args.fs,
+        min_overlap_sec=args.min_overlap, graph_k=args.k, sigma=args.sigma, radius=args.radius)