File size: 5,599 Bytes
9a5439c | 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 96 97 98 99 100 101 102 103 104 105 106 | """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): # = 33_
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): # = 33_
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): # = 33_
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): # = 33_ train_text_head
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())
|