File size: 5,386 Bytes
e8fa066 | 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 | """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))}
# collect test embeddings: (split, fold) -> {"if": {enc: z}, "markers": {enc: z}, "he": z, "tile_id": ids}
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) # our own files (contain numpy tile ids)
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())
|