Blame for experiments/20_pairing.py

"""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())