File size: 5,217 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 | """Emergent text alignment on ORION. The text head was trained ONLY on PathGen H&E-caption pairs (12_), the IF /
marker encoders ONLY on ORION H&E pairs (09_, --if_enc convnext_tiny --mk_enc mlp). Text and IF/markers never
co-occur in training; any text<->IF / text<->markers agreement is emergent through the shared UNI2 H&E space.
(A) concept probes: prompts.json (PathGen-style sentences) scored against every ORION test tile embedding
(he / if / markers); ground truth = marker-positive fraction (or nuclei count). Metrics: Spearman rho and
AUROC (top quartile vs rest), on tissue tiles (>=20 nuclei; all tiles for 'cellularity').
Baselines: zero-shot PLIP and ImageBind (their own image/text towers, H&E only, no training).
(B) retrieval with templated ORION descriptions (07_) as text queries through the PathGen-trained head
(domain-shifted: template style != PathGen captions), tie-aware R@1 on 100-tile galleries.
Outputs results/text_emergent_probes.csv, results/text_emergent_retrieval.csv
"""
import glob
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
TXT = ["plip", "biomedclip", "imagebind"]
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
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]
pos = np.isin(tiles, sub)
s = scores[pos]
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))}
probe_rows, ret_rows = [], []
for f in sorted(glob.glob(os.path.join(RUNS, "bench", "*_uni2_convnext_tiny_mlp_*", "test_embeddings.pt"))):
run = os.path.basename(os.path.dirname(f))
split, fold = run.split("_fold")[0], int(run.split("_fold")[1].split("_")[0])
E = torch.load(f, weights_only=False)
tiles, Z = np.asarray(E["tile_id"]), E["z"]
for enc in TXT:
prm = torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{enc}.npy"))[:, 0].astype(np.float32))
orion_txt = torch.from_numpy(np.load(os.path.join(CACHE, f"text_{enc}.npy"))[tiles, 0].astype(np.float32))
head = load_head(enc, prm.shape[1], Z["he"].shape[1])
with torch.no_grad():
zp = F.normalize(head(prm), dim=-1)
zt = F.normalize(head(orion_txt), dim=-1)
for i, (key, pr) in enumerate(prompts.items()):
for m in ["he", "if", "markers"]:
sc = (Z[m] @ zp[i]).numpy()
probe_rows.append({"split": split, "fold": fold, "text_enc": enc, "concept": key, "modality": m,
"kind": "bound (PathGen text head)", **probe(sc, tiles, key, pr["label"])})
for m in ["he", "if", "markers"]:
for q, g, name in [(zt, Z[m], f"text->{m}"), (Z[m], zt, f"{m}->text")]:
ret_rows.append({"split": split, "fold": fold, "text_enc": enc, "pair": name, **metrics(q, g)})
# zero-shot baselines on H&E with the model's own towers (computed once per fold)
for fm in ["plip", "imagebind"]:
hz = F.normalize(torch.from_numpy(np.load(os.path.join(CACHE, f"he_{fm}.npy"))[tiles, 0].astype(np.float32)), dim=-1)
pz = F.normalize(torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{fm}.npy"))[:, 0].astype(np.float32)), dim=-1)
for i, (key, pr) in enumerate(prompts.items()):
probe_rows.append({"split": split, "fold": fold, "text_enc": fm, "concept": key, "modality": "he",
"kind": "zero-shot (own image tower)", **probe((hz @ pz[i]).numpy(), tiles, key, pr["label"])})
P = pd.DataFrame(probe_rows); R = pd.DataFrame(ret_rows)
P.to_csv(os.path.join(OUT, "text_emergent_probes.csv"), index=False)
R.to_csv(os.path.join(OUT, "text_emergent_retrieval.csv"), index=False)
pd.set_option("display.width", 250); pd.set_option("display.max_rows", 300)
print(P.groupby(["kind", "text_enc", "modality", "split"])[["auroc", "spearman"]].mean().round(3).unstack("split").to_string())
print(P.groupby(["kind", "text_enc", "modality", "concept"])["auroc"].mean().unstack("concept").round(2).to_string())
print(R.groupby(["text_enc", "pair", "split"])["R@1_g100"].mean().unstack("split").round(3).to_string())
|