File size: 3,736 Bytes
2188a91 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """ADFTD (OpenNeuro ds004504) preprocessing, following DL4mHealth/Medformer.
Steps (verbatim from Medformer's ADFTD_preprocessing.ipynb):
* read the ASR/ICA-cleaned derivatives (.set, 19 channels, 500 Hz)
* linear-interpolation resample 500 -> 256 Hz
* segment into non-overlapping 256-timestamp windows, centred on the recording
* label map {C: 0 healthy, F: 1 FTD, A: 2 AD}
Split (from Medformer's ADFTDLoader): subject-wise 60/20/20 within each class.
Normalisation: per-sample, per-channel standardisation (uea.normalize_batch_ts).
"""
import argparse
import os
import mne
import numpy as np
import pandas as pd
from scipy import interpolate
mne.set_log_level("ERROR")
LABEL_MAP = {"A": 2, "F": 1, "C": 0}
def resample(arr, freq=500, target=256):
t = np.linspace(1, len(arr), len(arr))
f = interpolate.interp1d(t, arr, kind="linear")
t_new = np.linspace(1, len(arr), int(len(arr) / freq * target))
return f(t_new)
def segment(mat, window=256):
"""mat: (T, C). Centre-anchored non-overlapping windows."""
res = []
start = mat.shape[0] // 2
left = start - (start // window) * window
right = start + ((mat.shape[0] - start) // window) * window
for i in range(left, right, window):
res.append(mat[i:i + window])
return res
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--root", default="data/ADFTD")
ap.add_argument("--out", default="data/ADFTD/processed")
ap.add_argument("--window", type=int, default=256)
args = ap.parse_args()
parts = pd.read_csv(os.path.join(args.root, "participants.tsv"), sep="\t")
labels = {}
for row in parts.values:
pid = int(str(row[0])[-3:])
labels[pid] = LABEL_MAP[row[3]]
deriv = os.path.join(args.root, "derivatives")
feats, ys, subs = [], [], []
for pid in sorted(labels):
path = os.path.join(deriv, f"sub-{pid:03d}", "eeg",
f"sub-{pid:03d}_task-eyesclosed_eeg.set")
raw = mne.io.read_raw_eeglab(path, preload=True)
sig = raw.get_data() # (C, T) at 500 Hz
res = np.stack([resample(sig[c]) for c in range(sig.shape[0])], axis=1) # (T', C)
for win in segment(res, args.window):
feats.append(win.astype(np.float32))
ys.append(labels[pid])
subs.append(pid)
print(f"sub-{pid:03d} label={labels[pid]} raw={sig.shape} "
f"resampled={res.shape} windows={len(segment(res, args.window))}", flush=True)
X = np.stack(feats) # (N, 256, 19)
y = np.array(ys, dtype=np.int64)
s = np.array(subs, dtype=np.int64)
print("total", X.shape, np.bincount(y))
# subject-wise 60/20/20 within each class (Medformer ADFTDLoader, a=0.6, b=0.8)
train_ids, val_ids, test_ids = [], [], []
for cls in (0, 1, 2):
ids = sorted([p for p in labels if labels[p] == cls])
a, b = int(0.6 * len(ids)), int(0.8 * len(ids))
train_ids += ids[:a]
val_ids += ids[a:b]
test_ids += ids[b:]
print("subjects train/val/test:", len(train_ids), len(val_ids), len(test_ids))
os.makedirs(args.out, exist_ok=True)
for name, ids in (("train", train_ids), ("val", val_ids), ("test", test_ids)):
m = np.isin(s, ids)
Xi = X[m]
mu = Xi.mean(axis=1, keepdims=True)
sd = Xi.std(axis=1, keepdims=True)
sd[sd == 0] = 1.0
Xi = ((Xi - mu) / sd).astype(np.float32)
np.save(os.path.join(args.out, f"X_{name}.npy"), Xi)
np.save(os.path.join(args.out, f"y_{name}.npy"), y[m])
np.save(os.path.join(args.out, f"subj_{name}.npy"), s[m])
print(name, Xi.shape, np.bincount(y[m]))
if __name__ == "__main__":
main()
|