| """Re-check of the H&E <-> text encoder comparison (20_pairing.py, 28_) under the final protocol with CLEAN real text: |
| TCGA splits/final/tcga_fold<k>.csv (patient 5-fold, val = early stopping only), junk texts removed (40_clean_text.py), |
| identical text_head recipe to 33_final_eval.py (frozen image anchor, MLP on frozen text features, multi-positive InfoNCE). |
| Image anchors x text towers; PathGen-CLIP is left out (trained on PathGen-1.6M = most of this pool -> not held out). |
| Output results/text_enc_clean_long.csv, text_enc_clean_summary.csv |
| """ |
| import os |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from scipy.optimize import linear_sum_assignment |
|
|
| from common import CACHE, ROOT |
|
|
| IMGS = ["uni2", "phikonv2", "gigapath", "conch"] |
| TXTS = ["biomedclip", "plip", "imagebind", "conchclip"] |
| K, G100, DRAWS = 5, 100, 50 |
| dev = "cuda" |
| T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv")) |
| t_img = torch.as_tensor(T.img_idx.values, device=dev) |
| T = T[~pd.read_csv(os.path.join(CACHE, "alltxt_junk.csv")).junk.values] |
|
|
|
|
| def mlp(din, dout, hid=512): |
| return nn.Sequential(nn.Linear(din, hid), nn.GELU(), nn.Dropout(0.1), nn.Linear(hid, hid), nn.GELU(), |
| nn.LayerNorm(hid), nn.Linear(hid, dout)) |
|
|
|
|
| def one_text_per_image(t_ids, rng): |
| pick = T.loc[t_ids].groupby("img_idx").apply(lambda g: g.index[rng.integers(len(g))], include_groups=False) |
| return pick.index.values, pick.values |
|
|
|
|
| def g100(a, b, seed=0): |
| rng = np.random.default_rng(seed) |
| out = {m: [] for m in ["R@1", "R@5", "R@10", "MRR", "pair_acc"]} |
| for _ in range(DRAWS): |
| s = torch.from_numpy(rng.choice(len(a), min(G100, len(a)), replace=False)) |
| sim = (a[s] @ b[s].T).numpy(); d = np.diag(sim) |
| for S in (sim, sim.T): |
| rank = (S > d[:, None]).sum(1) |
| for k_ in (1, 5, 10): |
| out[f"R@{k_}"].append((rank < k_).mean() / 2) |
| out["MRR"].append((1 / (rank + 1)).mean() / 2) |
| ri, ci = linear_sum_assignment(-sim) |
| out["pair_acc"].append((ci[np.argsort(ri)] == np.arange(len(s))).mean()) |
| return {m: float(np.sum(v) / DRAWS) for m, v in out.items()} |
|
|
|
|
| def train(X, Tx, tr_t, va_t, epochs=60, bs=256, lr=3e-4, seed=0): |
| torch.manual_seed(seed); rng = np.random.default_rng(seed) |
| th = mlp(Tx.shape[1], X.shape[-1]).to(dev) |
| logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32, device=dev)) |
| opt = torch.optim.AdamW(list(th.parameters()) + [logit], lr=lr, weight_decay=0.05) |
| sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=lr, total_steps=epochs * max(1, len(tr_t) // bs), pct_start=0.05) |
| va_i, va_tt = one_text_per_image(va_t, np.random.default_rng(0)) |
| za_va = F.normalize(X[torch.as_tensor(va_i, device=dev), 0], dim=-1) |
| best, state = -1, None |
| for ep in range(epochs): |
| th.train() |
| perm = rng.permutation(tr_t) |
| for s in range(0, len(perm) - bs + 1, bs): |
| b = torch.as_tensor(perm[s:s + bs], device=dev) |
| ii = t_img[b] |
| zi = F.normalize(X[ii, torch.randint(0, X.shape[1], (len(b),), device=dev)], dim=-1) |
| lg = logit.exp().clamp(max=100) * zi @ F.normalize(th(Tx[b]), dim=-1).T |
| pos = (ii[:, None] == ii[None, :]).float() |
| tgt = pos / pos.sum(1, keepdim=True) |
| loss = (torch.sum(-tgt * F.log_softmax(lg, 1), 1).mean() + torch.sum(-tgt * F.log_softmax(lg.T, 1), 1).mean()) / 2 |
| opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sched.step() |
| th.eval() |
| with torch.no_grad(): |
| sim = za_va @ F.normalize(th(Tx[torch.as_tensor(va_tt, device=dev)]), dim=-1).T |
| v = float((((sim > sim.diag()[:, None]).sum(1) == 0).float().mean() + ((sim.T > sim.diag()[:, None]).sum(1) == 0).float().mean()) / 2) |
| if v > best: |
| best, state = v, {k_: x.detach().clone() for k_, x in th.state_dict().items()} |
| th.load_state_dict(state); th.eval() |
| return th, best |
|
|
|
|
| rows = [] |
| for ie in IMGS: |
| X = torch.from_numpy(np.load(os.path.join(CACHE, f"allimg_he_{ie}.npy")).astype(np.float32)).to(dev) |
| for te_ in TXTS: |
| Tx = torch.from_numpy(np.load(os.path.join(CACHE, f"alltxt_text_{te_}.npy"))[:, 0].astype(np.float32)).to(dev) |
| for k in range(K): |
| sp = pd.read_csv(os.path.join(ROOT, "splits", "final", f"tcga_fold{k}.csv")).set_index("img_idx").split |
| tsp = sp.loc[T.img_idx].values |
| tr_t, va_t, te_t = (T.index[tsp == x].values for x in ["train", "val", "test"]) |
| th, best = train(X, Tx, tr_t, va_t) |
| ii, tt = one_text_per_image(te_t, np.random.default_rng(k)) |
| with torch.no_grad(): |
| za = F.normalize(X[torch.as_tensor(ii, device=dev), 0], dim=-1).cpu() |
| zt = F.normalize(th(Tx[torch.as_tensor(tt, device=dev)]), dim=-1).cpu() |
| rows.append({"img_enc": ie, "txt_enc": te_, "fold": k, "val_R@1": best, "n_test": len(ii), **g100(za, zt)}) |
| print(f"{ie} + {te_} fold {k}: val {best:.3f} test R@1 {rows[-1]['R@1']:.3f}", flush=True) |
| R = pd.DataFrame(rows) |
| R.to_csv(os.path.join(ROOT, "results", "text_enc_clean_long.csv"), index=False) |
| S = R.groupby(["img_enc", "txt_enc"], sort=False)[["R@1", "R@10", "pair_acc"]].agg(["mean", "std"]).round(4) |
| S.to_csv(os.path.join(ROOT, "results", "text_enc_clean_summary.csv")) |
| pd.set_option("display.width", 200) |
| print(R.pivot_table(index="img_enc", columns="txt_enc", values="R@1", sort=False).round(3).to_string()) |
|
|
