| """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 |
| from embed import MIF_MARKERS, RNA_E, TXT_E, Embedder |
|
|
| 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)) |
|
|
|
|
| |
| 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() |
|
|
| |
| T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv")) |
| txt_ids = TI["tcga_txt"][tt_img] |
| 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() |
|
|
| |
| from rna_sets import load |
|
|
| 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() |
|
|
|
|
| |
| 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() |
|
|
| |
| 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]") |
|
|
