| """Evaluate modality-specific encoders: from-scratch baselines (runs/bench, 09_bench.py) vs frozen modality-specific |
| foundation models + MLP head (runs/bench_fm, 17_bench_fm.py), all bound to the frozen UNI2 H&E anchor. |
| |
| (A) retrieval, tie-aware R@1 on 100-tile galleries + Hungarian pairing acc on the full test set: |
| trained pairs he<->if, he<->markers (per encoder) |
| emergent pairs if<->markers (never trained together) (every if_enc x mk_enc combination of a fold) |
| (B) emergent text: the text head trained ONLY on PathGen H&E-caption pairs (12_train_text_pathgen.py, one per text |
| encoder) scores the concept prompts (prompts.json) against ORION he / if / markers embeddings; Spearman + AUROC |
| (same `probe` as 13_eval_text_emergent.py). Text never co-occurs with IF or markers in training. |
| Outputs results/fm_retrieval.csv, results/fm_text_probes.csv |
| """ |
| import glob |
| import itertools |
| import json |
| import os |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import torch.nn.functional as F |
| from scipy.stats import spearmanr |
| from sklearn.metrics import roc_auc_score |
|
|
| from common import CACHE, ROOT, RUNS, load_all_tiles, metrics |
|
|
| OUT = os.path.join(ROOT, "results") |
| prompts = json.load(open(os.path.join(ROOT, "prompts.json"))) |
| df = load_all_tiles() |
| stats = pd.read_csv(os.path.join(OUT, "tile_stats.csv")).set_index("tile_id") |
| df["n_nuclei"] = stats.loc[df.tile_id, "n_nuclei"].values |
| TXT = sorted(os.path.basename(d).split("_", 1)[1] for d in glob.glob(os.path.join(RUNS, "text_pathgen", "uni2_*")) |
| if os.path.exists(os.path.join(d, "head.pt")) and os.path.exists(os.path.join(CACHE, f"prompts_text_{os.path.basename(d).split('_', 1)[1]}.npy"))) |
|
|
|
|
| def load_head(enc, din, dout): |
| h = torch.nn.Sequential(torch.nn.Linear(din, 512), torch.nn.GELU(), torch.nn.Dropout(0.1), torch.nn.Linear(512, 512), |
| torch.nn.GELU(), torch.nn.LayerNorm(512), torch.nn.Linear(512, dout)) |
| h.load_state_dict(torch.load(os.path.join(RUNS, "text_pathgen", f"uni2_{enc}", "head.pt"), map_location="cpu")) |
| return h.eval() |
|
|
|
|
| def probe(scores, tiles, key, label_col): |
| sub = tiles if key == "cellularity" else tiles[df.loc[tiles, "n_nuclei"].values >= 20] |
| s = scores[np.isin(tiles, sub)] |
| y = df.loc[sub, label_col].values.astype(float) |
| if label_col != "n_nuclei": |
| y = y / np.maximum(df.loc[sub, "n_nuclei"].values, 1) |
| lab = y >= np.quantile(y, 0.75) |
| return {"spearman": float(spearmanr(s, y).correlation), |
| "auroc": float(roc_auc_score(lab, s)) if 0 < lab.sum() < len(lab) else np.nan, "n": int(len(sub))} |
|
|
|
|
| |
| folds = {} |
| for root, src in [("bench", "scratch"), ("bench_fm", None)]: |
| for f in glob.glob(os.path.join(RUNS, root, "*", "test_embeddings.pt")): |
| a = json.load(open(os.path.join(os.path.dirname(f), "train.json")))["args"] |
| if a["anchor"] != "uni2": |
| continue |
| E = torch.load(f, weights_only=False) |
| d = folds.setdefault((a["split"], a["fold"]), {"if": {}, "markers": {}, "he": E["z"]["he"], "tile_id": np.asarray(E["tile_id"])}) |
| assert (d["tile_id"] == np.asarray(E["tile_id"])).all() |
| d["if"][a["if_enc"]] = E["z"]["if"] |
| d["markers"][a["mk_enc"]] = E["z"]["markers"] |
|
|
| ret, prb = [], [] |
| for (split, fold), d in sorted(folds.items()): |
| he, tiles = d["he"], d["tile_id"] |
| for m in ["if", "markers"]: |
| for e, z in d[m].items(): |
| for q, g, name in [(he, z, f"he->{m}"), (z, he, f"{m}->he")]: |
| ret.append({"split": split, "fold": fold, "kind": "trained", "pair": name, "enc": e, **metrics(q, g)}) |
| for ei, em in itertools.product(d["if"], d["markers"]): |
| zi, zm = d["if"][ei], d["markers"][em] |
| for q, g, name in [(zi, zm, "if->markers"), (zm, zi, "markers->if")]: |
| ret.append({"split": split, "fold": fold, "kind": "emergent", "pair": name, "enc": f"{ei} + {em}", **metrics(q, g)}) |
| for t in TXT: |
| prm = torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{t}.npy"))[:, 0].astype(np.float32)) |
| with torch.no_grad(): |
| zp = F.normalize(load_head(t, prm.shape[1], he.shape[1])(prm), dim=-1) |
| embs = [("he", "uni2", he)] + [(m, e, z) for m in ["if", "markers"] for e, z in d[m].items()] |
| for m, e, z in embs: |
| for i, (key, pr) in enumerate(prompts.items()): |
| prb.append({"split": split, "fold": fold, "text_enc": t, "modality": m, "enc": e, "concept": key, |
| **probe((z @ zp[i]).numpy(), tiles, key, pr["label"])}) |
|
|
| R, P = pd.DataFrame(ret), pd.DataFrame(prb) |
| R.to_csv(os.path.join(OUT, "fm_retrieval.csv"), index=False) |
| P.to_csv(os.path.join(OUT, "fm_text_probes.csv"), index=False) |
| pd.set_option("display.width", 250); pd.set_option("display.max_rows", 500) |
| R["pair2"] = R.pair.map(lambda p: "<->".join(sorted(p.split("->")))) |
| g = R.groupby(["kind", "pair2", "enc", "split", "fold"])[["R@1_g100", "pair_acc"]].mean().reset_index() |
| print(g.groupby(["kind", "pair2", "enc", "split"])[["R@1_g100", "pair_acc"]].agg(["mean", "std"]).round(3).to_string()) |
| print(P.groupby(["text_enc", "modality", "enc", "split"])[["auroc", "spearman"]].mean().round(3).unstack("split").to_string()) |
|
|
