File size: 4,978 Bytes
c2ede1f | 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 | """End-to-end check of embed.py from RAW inputs on examples/ (64 held-out fold-0 test items per dataset; needs the
frozen encoders, see setup_encoders.sh; GPU).
python -I example_check.py [models/ej_ad_mv_fold0]
1. raw input -> encoder features, compared with the cached features the model was trained / tested on
(examples/ref_feats.npz): per-item cosine. Expected ~1.000; Tahoe-x1 / SE-600M ~0.99 (their official input
pipelines sample genes at random); VirTues lower here because mIF standardisation statistics are computed over the
tiles given in one call (64 tiles here vs the whole ORION set for the cache).
2. raw input -> 1536-d embeddings -> retrieval inside the 64 items (R@1, both directions; chance 1/64), next to the
same retrieval from the cached features.
Output: examples/example_check.csv; exit code 1 if a deterministic encoder does not reproduce its cached features.
"""
import os
import sys
import time
import numpy as np
import pandas as pd
import scipy.sparse as sp
import torch
ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, ROOT)
from embed import RNA_E, TXT_E, Embedder # noqa: E402
MD = sys.argv[1] if len(sys.argv) > 1 else os.path.join(ROOT, "models", "ej_ad_mv_fold0")
EX = os.path.join(ROOT, "examples")
E = Embedder(MD)
R = dict(np.load(os.path.join(EX, "ref_feats.npz")))
rows, t0 = [], time.time()
MIN_COS = {"uni2": .999, "medvtok": .999, "keep": .999, "biomedclip": .99, "qwen3-0.6b": .999, "hvg": .999,
"aidocell": .999, "tahoex1": .97, "se600m": .97, "virtues": 0.0} # median cosine required (virtues: report only)
def cos(a, b):
a, b = torch.as_tensor(np.asarray(a, np.float32)), torch.as_tensor(np.asarray(b, np.float32))
return torch.nn.functional.cosine_similarity(a, b, dim=-1).numpy()
def feat(ds, enc, raw):
c = cos(raw, R[f"{ds}_{enc}"])
rows.append({"check": "features", "data": ds, "encoder": enc, "n": len(c), "cos_median": float(np.median(c)),
"cos_min": float(c.min()), "ok": bool(np.median(c) >= MIN_COS[enc])})
print(f"[{time.time() - t0:4.0f}s] {ds:5s} {enc:11s} cos median {np.median(c):.4f} min {c.min():.4f}"
f"{'' if rows[-1]['ok'] else ' <-- FAIL'}", flush=True)
def r1(a, b):
s = a @ b.T; d = s.diag()
return float((((s > d[:, None]).sum(1) == 0).float().mean() + ((s.T > d[:, None]).sum(1) == 0).float().mean()) / 2)
# raw H&E (PNG files / uint8 arrays) -> UNI2 + MedVTok
U, M = {}, {}
U["orion"], M["orion"] = E.he_feats(sorted(os.path.join(EX, "orion_he", f) for f in os.listdir(os.path.join(EX, "orion_he"))))
U["hest"], M["hest"] = E.he_feats(np.load(os.path.join(EX, "hest_he.npy")))
C = pd.read_csv(os.path.join(EX, "tcga_captions.csv"))
U["tcga"], M["tcga"] = E.he_feats([os.path.join(EX, "tcga_he", f) for f in C.file])
for ds in U:
feat(ds, "uni2", U[ds]); feat(ds, "medvtok", M[ds])
E.release()
# raw text -> KEEP / BiomedCLIP / Qwen3
Ft = E.text_feats(C.caption.fillna("").astype(str).tolist())
for e in TXT_E:
feat("tcga", e, Ft[e])
E.release()
# raw counts -> hvg / Tahoe-x1 / SE-600M / AIDO.Cell
genes = pd.read_csv(os.path.join(EX, "hest_genes.csv")).gene.values
Fr = E.rna_feats(sp.load_npz(os.path.join(EX, "hest_counts.npz")), genes)
for e in RNA_E:
feat("hest", e, Fr[e])
E.release()
# raw mIF tiles -> VirTues
m = np.load(os.path.join(EX, "orion_mif.npz"), allow_pickle=True)
Fm = E.mif_feats(m["tiles"], m["slide"])
feat("orion", "virtues", Fm)
E.release()
# embeddings + retrieval inside the 64 items
pairs = {"H&E<->mIF (ORION)": ((E.he_from_feats(U["orion"], M["orion"]), E.mif_from_feats(Fm)),
(E.he_from_feats(R["orion_uni2"], R["orion_medvtok"]), E.mif_from_feats(R["orion_virtues"]))),
"H&E<->RNA (HEST)": ((E.he_from_feats(U["hest"], M["hest"]), E.rna_from_feats(Fr)),
(E.he_from_feats(R["hest_uni2"], R["hest_medvtok"]),
E.rna_from_feats({e: R[f"hest_{e}"] for e in RNA_E}))),
"H&E<->text (TCGA)": ((E.he_from_feats(U["tcga"], M["tcga"]), E.text_from_feats(Ft)),
(E.he_from_feats(R["tcga_uni2"], R["tcga_medvtok"]),
E.text_from_feats({e: R[f"tcga_{e}"] for e in TXT_E})))}
for p, ((a, b), (ac, bc)) in pairs.items():
rows.append({"check": "retrieval", "data": p, "n": len(a), "R@1_raw": r1(a, b), "R@1_cached": r1(ac, bc),
"chance": 1 / len(a)})
print(f"{p:18s} R@1 within {len(a)}: raw inputs {rows[-1]['R@1_raw']:.3f} | cached features "
f"{rows[-1]['R@1_cached']:.3f} (chance {1 / len(a):.3f})", flush=True)
D = pd.DataFrame(rows)
D.to_csv(os.path.join(EX, "example_check.csv"), index=False)
bad = D[(D.check == "features") & ~D.ok.astype(bool)]
print(f"done [{time.time() - t0:.0f}s]:", "all encoders reproduce the cached features" if bad.empty else
f"FAILED {bad[['data', 'encoder']].values.tolist()}")
sys.exit(0 if bad.empty else 1)
|