File size: 8,059 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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """Check embed.py from RAW inputs (pixels, strings, counts) against the cached training features and the fold-0 test.
.venv/bin/python -I 53_embed_check.py [n] (n = items per raw check, default 300)
1. feature reproduction: embed.py raw-input encoders vs the cached features the model was trained / tested on
(per-item cosine of the raw encoder features), for every encoder:
H&E UNI2 + MedVTok ORION tiles (PNG), TCGA/PathGen patches (PNG), HEST spot crops (cache/hest/he_img.npy)
text KEEP / BiomedCLIP / Qwen3 TCGA real captions (alltxt)
RNA hvg / Tahoe-x1 / SE-600M / AIDO.Cell HEST spots (rna_sets.load('hest') rows)
mIF VirTues (all 3187 ORION tiles: standardisation statistics are per slide / per dataset)
2. end-to-end: models/ej_ad_mv_fold0 on held-out fold-0 test items, embeddings from raw inputs vs from cached
features (same items, g100 R@1)
Output results/embed_check.csv
"""
import os
import sys
import time
import numpy as np
import pandas as pd
import torch
ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, ROOT)
from common import CACHE, ORION, load_all_tiles # noqa: E402
from embed import MIF_MARKERS, RNA_E, TXT_E, Embedder # noqa: E402
N = int(sys.argv[1]) if len(sys.argv) > 1 else 300
MD = os.path.join(ROOT, "models", "ej_ad_mv_fold0")
E = Embedder(MD)
TI = np.load(os.path.join(MD, "test_items.npz"))
rng = np.random.default_rng(0)
rows, t0 = [], time.time()
f32 = lambda a: np.asarray(a, np.float32)
def cos(a, b):
a, b = torch.as_tensor(f32(a)), torch.as_tensor(f32(b))
return torch.nn.functional.cosine_similarity(a, b, dim=-1).numpy()
def report(mod, enc, src, c):
rows.append({"check": "features", "modality": mod, "encoder": enc, "source": src, "n": len(c),
"cos_median": float(np.median(c)), "cos_min": float(c.min())})
print(f"[{time.time() - t0:.0f}s] {mod:5s} {enc:11s} {src:14s} n={len(c)} cos median {np.median(c):.4f} min {c.min():.4f}",
flush=True)
def r1_g100(a, b, seed=0):
a, b = torch.nn.functional.normalize(a, dim=-1), torch.nn.functional.normalize(b, dim=-1)
g = np.random.default_rng(seed)
out = []
for _ in range(50):
s = torch.from_numpy(g.choice(len(a), min(100, len(a)), replace=False))
S = a[s] @ b[s].T; d = S.diag()
out.append((((S > d[:, None]).sum(1) == 0).float().mean() + ((S.T > d[:, None]).sum(1) == 0).float().mean()).item() / 2)
return float(np.mean(out))
# ------------------------------------------------------------------ H&E
tiles = load_all_tiles()
ot = rng.choice(TI["orion"], min(N, len(TI["orion"])), replace=False)
U, M = E.he_feats([os.path.join(ORION, p) for p in tiles.image_path.values[ot]])
report("H&E", "uni2", "ORION", cos(U, np.load(os.path.join(CACHE, "he_uni2.npy"), mmap_mode="r")[ot, 0]))
report("H&E", "medvtok", "ORION", cos(M, np.load(os.path.join(CACHE, "he_medvtok.npy"), mmap_mode="r")[ot, 0]))
he_raw = {"orion": (ot, U, M)}
I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
tt_img = rng.choice(len(TI["tcga_img"]), min(N, len(TI["tcga_img"])), replace=False)
ti = TI["tcga_img"][tt_img]
U, M = E.he_feats(I.image_path.values[ti])
report("H&E", "uni2", "TCGA/PathGen", cos(U, np.load(os.path.join(CACHE, "allimg_he_uni2.npy"), mmap_mode="r")[ti, 0]))
report("H&E", "medvtok", "TCGA/PathGen", cos(M, np.load(os.path.join(CACHE, "allimg_he_medvtok.npy"), mmap_mode="r")[ti, 0]))
he_raw["tcga"] = (tt_img, U, M)
hsp = rng.choice(TI["hest"], min(N, len(TI["hest"])), replace=False)
U, M = E.he_feats(np.load(os.path.join(CACHE, "hest", "he_img.npy"), mmap_mode="r")[hsp])
report("H&E", "uni2", "HEST", cos(U, np.load(os.path.join(CACHE, "hest", "he_uni2.npy"), mmap_mode="r")[hsp, 0]))
report("H&E", "medvtok", "HEST", cos(M, np.load(os.path.join(CACHE, "hest", "he_medvtok.npy"), mmap_mode="r")[hsp, 0]))
he_raw["hest"] = (hsp, U, M)
E.release()
# ------------------------------------------------------------------ text
T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv"))
txt_ids = TI["tcga_txt"][tt_img] # the captions paired with the TCGA test images above
Ft = E.text_feats(T.caption.fillna("").astype(str).values[txt_ids].tolist())
for e in TXT_E:
report("text", e, "TCGA captions", cos(Ft[e], np.load(os.path.join(CACHE, f"alltxt_text_{e}.npy"), mmap_mode="r")[txt_ids, 0]))
E.release()
# ------------------------------------------------------------------ RNA
from rna_sets import load # noqa: E402
Xh, genes = load("hest")
Fr = E.rna_feats(Xh[hsp], genes)
for e in RNA_E:
ref = (np.load(os.path.join(CACHE, "hest", "hvg.npy"), mmap_mode="r") if e == "hvg" else
np.load(os.path.join(CACHE, "rnafm", f"{e}_hest.npy"), mmap_mode="r"))[hsp]
report("RNA", e, "HEST", cos(Fr[e], ref))
E.release()
# ------------------------------------------------------------------ mIF (whole ORION set: dataset-level standardisation)
class LazyIF:
"""(N, 16, 333, 333) ORION mIF tiles read on demand (uint8 tifs, 17 channels -> MIF_MARKERS)"""
def __init__(self, df):
import tifffile
self.paths, self.rd = [os.path.join(ORION, p) for p in df.target_path], tifffile.imread
V = E._get("virtues")[0]
names = [m for m, _ in V.MARKERS]
self.ch = [names.index(m) for m in MIF_MARKERS]
def __len__(self):
return len(self.paths)
def __getitem__(self, s):
return np.stack([self.rd(p)[..., self.ch].transpose(2, 0, 1) for p in self.paths[s]])
Fm = E.mif_feats(LazyIF(tiles), tiles.in_slide_name.values)
report("mIF", "virtues", "ORION (all)", cos(Fm, np.load(os.path.join(CACHE, "if_virtues_m16.npy"), mmap_mode="r")[:, 0]))
E.release()
# ------------------------------------------------------------------ end-to-end on held-out fold-0 test items
UO, MO = np.load(os.path.join(CACHE, "he_uni2.npy"), mmap_mode="r"), np.load(os.path.join(CACHE, "he_medvtok.npy"), mmap_mode="r")
IFc = np.load(os.path.join(CACHE, "if_virtues_m16.npy"), mmap_mode="r")
ot, U, M = he_raw["orion"]
e2e = {"H&E<->mIF (ORION)": ((E.he_from_feats(U, M), E.mif_from_feats(Fm[ot])),
(E.he_from_feats(UO[ot, 0], MO[ot, 0]), E.mif_from_feats(IFc[ot, 0])))}
hsp, U, M = he_raw["hest"]
ref_r = {e: (np.load(os.path.join(CACHE, "hest", "hvg.npy"), mmap_mode="r") if e == "hvg" else
np.load(os.path.join(CACHE, "rnafm", f"{e}_hest.npy"), mmap_mode="r"))[hsp] for e in RNA_E}
e2e["H&E<->RNA (HEST)"] = ((E.he_from_feats(U, M), E.rna_from_feats(Fr)),
(E.he_from_feats(np.load(os.path.join(CACHE, "hest", "he_uni2.npy"), mmap_mode="r")[hsp, 0],
np.load(os.path.join(CACHE, "hest", "he_medvtok.npy"), mmap_mode="r")[hsp, 0]),
E.rna_from_feats(ref_r)))
_, U, M = he_raw["tcga"]
ref_t = {e: np.load(os.path.join(CACHE, f"alltxt_text_{e}.npy"), mmap_mode="r")[txt_ids, 0] for e in TXT_E}
e2e["H&E<->text (TCGA)"] = ((E.he_from_feats(U, M), E.text_from_feats(Ft)),
(E.he_from_feats(np.load(os.path.join(CACHE, "allimg_he_uni2.npy"), mmap_mode="r")[ti, 0],
np.load(os.path.join(CACHE, "allimg_he_medvtok.npy"), mmap_mode="r")[ti, 0]),
E.text_from_feats(ref_t)))
for pair, v in e2e.items():
(a_raw, b_raw), (a_c, b_c) = v
rows.append({"check": "end-to-end", "modality": pair, "n": len(a_raw), "R@1_raw": r1_g100(a_raw, b_raw),
"R@1_cached": r1_g100(a_c, b_c), "emb_cos_median_a": float(np.median(cos(a_raw, a_c))),
"emb_cos_median_b": float(np.median(cos(b_raw, b_c)))})
print(f"{pair:20s} n={len(a_raw)} g100 R@1 raw-input {rows[-1]['R@1_raw']:.3f} | cached features "
f"{rows[-1]['R@1_cached']:.3f}; embedding cos median {rows[-1]['emb_cos_median_a']:.4f} / "
f"{rows[-1]['emb_cos_median_b']:.4f}", flush=True)
pd.DataFrame(rows).to_csv(os.path.join(ROOT, "results", "embed_check.csv"), index=False)
print(f"done [{time.time() - t0:.0f}s]")
|