"""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]")