File size: 4,151 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 | """PTB-XL preprocessing, following DL4mHealth/Medformer's PTB-XL_preprocessing.ipynb.
* records500 (12-lead, 500 Hz, 10 s) -> linear resample to 250 Hz -> 2500 samples
* per-record StandardScaler over time, then reshape into 10 non-overlapping
250-timestamp windows (the paper's Table 5: T=250, patch 10, 5 classes)
* label = the SCP code with the highest likelihood, mapped to the 5 diagnostic
superclasses NORM / MI / STTC / CD / HYP
* subject-wise (patient_id) split; Medformer's PTB-XL loader uses a 60/20/20
patient split, which we reproduce with a fixed seed.
"""
import argparse
import ast
import os
import numpy as np
import pandas as pd
import wfdb
from scipy import interpolate
from sklearn.preprocessing import StandardScaler
SUPER = {"NORM": 0, "MI": 1, "STTC": 2, "CD": 3, "HYP": 4}
def build_scp_map(scp_csv):
"""diagnostic_class column of scp_statements.csv maps each SCP code to a superclass."""
df = pd.read_csv(scp_csv, index_col=0)
df = df[df.diagnostic == 1]
return {code: row.diagnostic_class for code, row in df.iterrows()
if isinstance(row.diagnostic_class, str) and row.diagnostic_class in SUPER}
def resample(arr, freq=500, target=250):
t = np.linspace(1, len(arr), len(arr))
f = interpolate.interp1d(t, arr, kind="linear")
return f(np.linspace(1, len(arr), int(len(arr) / freq * target)))
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--root", default="data/PTBXL")
ap.add_argument("--out", default="data/PTBXL/processed")
ap.add_argument("--subject-frac", type=float, default=1.0,
help="fraction of patients to keep (compute budget)")
ap.add_argument("--seed", type=int, default=0)
a = ap.parse_args()
info = pd.read_csv(os.path.join(a.root, "ptbxl_database.csv"))
code2super = build_scp_map(os.path.join(a.root, "scp_statements.csv"))
rows = []
for r in info.itertuples():
codes = ast.literal_eval(r.scp_codes)
diag = {c: v for c, v in codes.items() if c in code2super}
if not diag:
continue # Medformer keeps only diagnosable records
best = max(diag, key=diag.get) # highest-likelihood SCP code
rows.append((r.ecg_id, r.patient_id, SUPER[code2super[best]], r.filename_hr))
df = pd.DataFrame(rows, columns=["ecg_id", "patient_id", "label", "path"])
print(f"records with a diagnostic superclass: {len(df)} "
f"({df.patient_id.nunique()} patients)")
print("class counts:", df.label.value_counts().sort_index().to_dict())
pats = np.array(sorted(df.patient_id.unique()))
rng = np.random.default_rng(a.seed)
rng.shuffle(pats)
if a.subject_frac < 1.0:
pats = pats[: max(2, int(len(pats) * a.subject_frac))]
df = df[df.patient_id.isin(pats)]
print(f"subsampled to {len(pats)} patients / {len(df)} records "
f"(subject_frac={a.subject_frac})")
n = len(pats)
splits = {"train": set(pats[: int(0.6 * n)]),
"val": set(pats[int(0.6 * n):int(0.8 * n)]),
"test": set(pats[int(0.8 * n):])}
os.makedirs(a.out, exist_ok=True)
for name, ids in splits.items():
sub = df[df.patient_id.isin(ids)]
X, y = [], []
for i, r in enumerate(sub.itertuples()):
sig, _ = wfdb.rdsamp(os.path.join(a.root, r.path)) # (5000, 12) @500 Hz
res = np.stack([resample(sig[:, c]) for c in range(sig.shape[1])], 1)
res = StandardScaler().fit_transform(res) # per-record z-score
w = res[: (res.shape[0] // 250) * 250].reshape(-1, 250, res.shape[1])
X.append(w.astype(np.float32))
y.append(np.full(w.shape[0], r.label, dtype=np.int64))
if i % 2000 == 0:
print(f" {name} {i}/{len(sub)}", flush=True)
X = np.concatenate(X)
y = np.concatenate(y)
np.save(os.path.join(a.out, f"X_{name}.npy"), X)
np.save(os.path.join(a.out, f"y_{name}.npy"), y)
print(name, X.shape, np.bincount(y, minlength=5), flush=True)
if __name__ == "__main__":
main()
|