"""Text <-> H&E image pairing benchmark on ALL real image-text pairs (unified pool from 19_pairsets_prep.py: 6144 TCGA patches, 6480 pairs; text types caption / summary / qa_open / qa_close). Protocol: patient-level 5-fold CV over the pool (inside each training fold 10% of patients are held out for early stopping). Training uses every (image, text) pair of the training patients; an image with several texts counts all of them as positives (multi-positive InfoNCE). Every image encoder x text encoder combination in three modes: zeroshot : the CLIP model's own image and text towers, no training (only when image enc == text enc) text_head : frozen image features = anchor; trainable MLP head on the frozen text features (ImageBind recipe) two_heads : trainable MLP heads on both frozen towers into a shared 512-d space (CLIP-style projection) Evaluation on the held-out fold, per text type and for "all" (one text per image): tie-aware R@1/5/10, MRR, Hungarian one-to-one pairing accuracy over the whole gallery, R@1 on random 100-pair galleries; both directions. Caveat: PathGen-CLIP-L was trained on PathGen-1.6M (most of this pool) -> its numbers are not held-out. Outputs results/pairing_long.csv, results/pairing_summary.csv """ import argparse import os import numpy as np import pandas as pd import torch import torch.nn as nn import torch.nn.functional as F from sklearn.model_selection import GroupKFold from common import CACHE, ROOT, metrics p = argparse.ArgumentParser() p.add_argument("--img", nargs="+", default=["uni2", "phikonv2", "plip", "biomedclip", "imagebind", "pathgenclip"]) p.add_argument("--txt", nargs="+", default=["plip", "biomedclip", "imagebind", "pathgenclip"]) p.add_argument("--epochs", type=int, default=60) p.add_argument("--bs", type=int, default=256) p.add_argument("--lr", type=float, default=3e-4) p.add_argument("--tag", default="") args = p.parse_args() dev = "cuda" OUT = os.path.join(ROOT, "results") I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv")) T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv")) TTYPES = ["caption", "summary", "qa_open", "qa_close"] def load(kind, enc): f = os.path.join(CACHE, f"all{kind}_{enc}.npy") return torch.from_numpy(np.load(f).astype(np.float32)).to(dev) if os.path.exists(f) else None IMG = {e: load("img_he", e) for e in args.img} # (n_img, 8, D) TXT = {e: (lambda a: a[:, 0] if a is not None else None)(load("txt_text", e)) for e in args.txt} # (n_txt, D) t_img = torch.as_tensor(T.img_idx.values, device=dev) def mlp(din, dout, hid=512): # same head as 12_train_text_pathgen.py / 09_bench.py 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 r1(a, b): s = a @ b.T return float(((s > s.diag()[:, None]).sum(1) == 0).float().mean()) def eval_pairs(t_ids, rng): """one text per image (random) -> aligned (img_ids, txt_ids) without duplicate images""" sub = T.loc[t_ids] pick = sub.groupby("img_idx").apply(lambda g: g.index[rng.integers(len(g))], include_groups=False).values return T.img_idx.values[pick], pick def train(ie, te, mode, tr_t, va_t, seed=0): torch.manual_seed(seed); rng = np.random.default_rng(seed) X, Tx = IMG[ie], TXT[te] th = mlp(Tx.shape[1], 512 if mode == "two_heads" else X.shape[-1]).to(dev) ih = mlp(X.shape[-1], 512).to(dev) if mode == "two_heads" else nn.Identity() logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32, device=dev)) opt = torch.optim.AdamW(list(th.parameters()) + list(ih.parameters()) + [logit], lr=args.lr, weight_decay=0.05) steps = args.epochs * max(1, len(tr_t) // args.bs) sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=steps, pct_start=0.05) va_i, va_tt = eval_pairs(va_t, np.random.default_rng(0)) def enc_img(ids, view=None): return F.normalize(ih(X[ids, 0 if view is None else view]), dim=-1) def enc_txt(ids): return F.normalize(th(Tx[ids]), dim=-1) best, state = -1, None for ep in range(args.epochs): th.train(); ih.train() perm = rng.permutation(tr_t) for s in range(0, len(perm) - args.bs + 1, args.bs): b = torch.as_tensor(perm[s:s + args.bs], device=dev) ii = t_img[b] zi = enc_img(ii, torch.randint(0, 8, (len(b),), device=dev)) zt = enc_txt(b) lg = logit.exp().clamp(max=100) * zi @ zt.T pos = (ii[:, None] == ii[None, :]).float() # same image -> positive (multi-positive InfoNCE) 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(); ih.eval() with torch.no_grad(): zi, zt = enc_img(torch.as_tensor(va_i, device=dev)), enc_txt(torch.as_tensor(va_tt, device=dev)) v = (r1(zi, zt) + r1(zt, zi)) / 2 if v > best: best = v state = ({k: x.clone() for k, x in th.state_dict().items()}, {k: x.clone() for k, x in ih.state_dict().items()}) th.load_state_dict(state[0]); ih.load_state_dict(state[1]); th.eval(); ih.eval() return (lambda i_ids, t_ids: (enc_img(torch.as_tensor(i_ids, device=dev)).cpu(), enc_txt(torch.as_tensor(t_ids, device=dev)).cpu())), best folds = list(GroupKFold(5).split(I, groups=I.patient)) rows = [] for k, (trv_img, te_img) in enumerate(folds): pats = np.random.default_rng(k).permutation(sorted(I.patient.iloc[trv_img].unique())) va_p = set(pats[: max(1, len(pats) // 10)]) img_split = pd.Series("test", index=I.index) img_split.iloc[trv_img] = np.where(I.patient.iloc[trv_img].isin(va_p), "val", "train") tsplit = img_split.loc[T.img_idx].values tr_t, va_t, te_t = (T.index[tsplit == s].values for s in ["train", "val", "test"]) # evaluation sets: per text type and 'all', one text per image ev = {"all": eval_pairs(te_t, np.random.default_rng(k))} for tt in TTYPES: ids = te_t[T.ttype.values[te_t] == tt] if len(ids) >= 20: ev[tt] = eval_pairs(ids, np.random.default_rng(k)) for ie in args.img: for te in args.txt: if IMG[ie] is None or TXT[te] is None: continue for mode in (["zeroshot"] if ie == te else []) + ["text_head", "two_heads"]: if mode == "zeroshot": f = lambda i_ids, t_ids: (F.normalize(IMG[ie][torch.as_tensor(i_ids, device=dev), 0], dim=-1).cpu(), F.normalize(TXT[te][torch.as_tensor(t_ids, device=dev)], dim=-1).cpu()) val = np.nan else: f, val = train(ie, te, mode, tr_t, va_t) for name, (i_ids, t_ids) in ev.items(): with torch.no_grad(): zi, zt = f(i_ids, t_ids) for q, g, d in [(zi, zt, "img->text"), (zt, zi, "text->img")]: rows.append({"eval": name, "fold": k, "img_enc": ie, "txt_enc": te, "mode": mode, "dir": d, "val_r1": val, **metrics(q, g)}) print(f"fold {k} {ie:11s} {te:11s} done (n_train_pairs={len(tr_t)}, n_test_img={len(ev['all'][0])})", flush=True) pd.DataFrame(rows).to_csv(os.path.join(OUT, f"pairing_long{args.tag}.csv"), index=False) R = pd.DataFrame(rows) cols = ["R@1", "R@5", "R@10", "pair_acc", "R@1_g100"] S = R.groupby(["eval", "mode", "img_enc", "txt_enc", "fold"])[cols + ["N"]].mean().groupby(["eval", "mode", "img_enc", "txt_enc"]).agg(["mean", "std"]) S.round(4).to_csv(os.path.join(OUT, f"pairing_summary{args.tag}.csv")) pd.set_option("display.width", 250); pd.set_option("display.max_rows", 500) for e in ["all"] + TTYPES: print(f"\n=== eval: {e} (mean over both directions and 5 folds) ===") print(R[R["eval"] == e].groupby(["mode", "img_enc", "txt_enc"])[cols + ["N"]].mean().round(3).sort_values("R@1_g100", ascending=False).to_string())