File size: 2,892 Bytes
7916811 | 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 | """Inference for DeepECG-SSL-CLSA-CVD (Model A / Model A+).
Requires: torch, safetensors, numpy, fairseq-signals (https://github.com/Jwoo5/fairseq-signals) and the public
DeepECG-SSL v2 (amplitude-preserved) pretrained encoder checkpoint, which defines the architecture.
ECG input: (N, 12, 2500) float32 in millivolts, 250 Hz, 10 s, standard lead order, no per-lead normalisation.
from inference import DeepECGSSLCVD
m = DeepECGSSLCVD("weights/", "path/to/deepecg-ssl-v2-wcr-amp-preserved-encoder.pt", device="cuda")
p_ecg = m.ecg_probability(ecg_mv) # Model A (calibrated ECG-only probability)
risk = m.model_a_plus(ecg_mv, age, female, hypertension) # Model A+ (no blood test)
"""
import json, os, numpy as np, torch, torch.nn as nn
from safetensors.torch import load_file
from fairseq_signals.utils import checkpoint_utils
HERE = os.path.dirname(os.path.abspath(__file__))
lg = lambda p: np.log(np.clip(p, 1e-6, 1 - 1e-6) / (1 - np.clip(p, 1e-6, 1 - 1e-6)))
class Net(nn.Module):
def __init__(self, base_ckpt):
super().__init__()
models, _, _ = checkpoint_utils.load_model_and_task(base_ckpt)
self.enc = models[0] if isinstance(models, list) else models
for n in ["quantizer", "project_q", "final_proj", "target_glu"]:
if hasattr(self.enc, n): setattr(self.enc, n, None)
self.head = nn.Sequential(nn.Dropout(0.3), nn.Linear(768, 1))
def forward(self, x):
return self.head(self.enc(source=x, mask=False, features_only=True)["x"].mean(1)).squeeze(1)
class DeepECGSSLCVD:
def __init__(self, weights_dir, base_ckpt, device="cpu"):
self.p = json.load(open(os.path.join(HERE, "model_params.json"))); self.device = device; self.models = []
for k in range(5):
m = Net(base_ckpt); m.load_state_dict(load_file(os.path.join(weights_dir, f"deepecg_ssl_clsa_cvd_fold{k}.safetensors")), strict=False)
self.models.append(m.to(device).eval())
@torch.no_grad()
def ecg_probability(self, ecg_mv, batch=64):
x = torch.as_tensor(np.asarray(ecg_mv, dtype=np.float32)); out = []
for k, m in enumerate(self.models):
pk = []
for i in range(0, len(x), batch):
pk.append(torch.sigmoid(m(x[i:i + batch].to(self.device))).float().cpu().numpy())
a, b = self.p["per_fold_platt"][str(k)]["a"], self.p["per_fold_platt"][str(k)]["b"]
out.append(1 / (1 + np.exp(-(a * lg(np.concatenate(pk)) + b))))
return np.mean(out, 0)
def model_a_plus(self, ecg_mv, age, female, hypertension):
c = self.p["model_A_plus"]; e = lg(self.ecg_probability(ecg_mv))
z = c["intercept"] + c["b_ecg"] * e + c["b_age"] * np.asarray(age, float) + c["b_female"] * np.asarray(female, float) + c["b_htn"] * np.asarray(hypertension, float)
return 1 / (1 + np.exp(-z))
|