| """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} |
| TXT = {e: (lambda a: a[:, 0] if a is not None else None)(load("txt_text", e)) for e in args.txt} |
| t_img = torch.as_tensor(T.img_idx.values, device=dev) |
|
|
|
|
| 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 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() |
| 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"]) |
| |
| 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()) |
|
|
