| """ImageBind-style binding of IF / nuclei / markers to a frozen H&E (ImageBind-Huge vision) anchor. |
| |
| loss = sum_m InfoNCE_sym(he, m) (--loss anchor, ImageBind default; IF<->nuclei etc. are *emergent*) |
| + sum_{m<n} InfoNCE_sym(m, n) (--loss allpairs ablation) |
| |
| Evaluation on the test split: cross-modal retrieval for all 12 directed modality pairs |
| (R@1/5/10, MRR, median rank) and one-to-one pairing accuracy (Hungarian assignment on similarity). |
| """ |
| import argparse |
| import itertools |
| import json |
| import os |
| import time |
|
|
| import numpy as np |
| import pandas as pd |
| import timm |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from scipy.optimize import linear_sum_assignment |
|
|
| from common import CACHE, MODALITIES, RUNS, SPLITS, dihedral |
|
|
| p = argparse.ArgumentParser() |
| p.add_argument("--split", default="orion_official") |
| p.add_argument("--fold", type=int, default=0) |
| p.add_argument("--loss", default="anchor", choices=["anchor", "allpairs"]) |
| p.add_argument("--anchor", default="imagebind", choices=["imagebind", "uni2", "conch", "gigapath", "plip", "phikonv2"], |
| help="frozen H&E encoder whose features (cache/he_<anchor>.npy) define the shared space") |
| p.add_argument("--epochs", type=int, default=40) |
| p.add_argument("--bs", type=int, default=128) |
| p.add_argument("--lr", type=float, default=3e-4) |
| p.add_argument("--backbone", default="resnet18.a1_in1k") |
| p.add_argument("--seed", type=int, default=0) |
| p.add_argument("--min_nuclei_tissue", type=int, default=20) |
| args = p.parse_args() |
| torch.manual_seed(args.seed); np.random.seed(args.seed) |
| dev = "cuda" |
| run_dir = os.path.join(RUNS, f"{args.split}_fold{args.fold}_{args.loss}" |
| + ("" if args.anchor == "imagebind" else f"_{args.anchor}")) |
| os.makedirs(run_dir, exist_ok=True) |
|
|
| |
| sp = pd.read_csv(os.path.join(SPLITS, args.split, f"fold{args.fold}.csv")) |
| he_emb = torch.from_numpy(np.load(os.path.join(CACHE, f"he_{args.anchor}.npy")).astype(np.float32)) |
| D = he_emb.shape[-1] |
| IF = torch.from_numpy(np.load(os.path.join(CACHE, "if.npy"))) |
| NUC = torch.from_numpy(np.load(os.path.join(CACHE, "nuclei.npy"))) |
| MK = torch.from_numpy(np.load(os.path.join(CACHE, "markers.npy"))) |
| tr_idx = sp.tile_id[sp.split == "train"].values |
| va_idx = sp.tile_id[sp.split == "val"].values |
| te_idx = sp.tile_id[sp.split == "test"].values |
| |
| if_mean = IF[tr_idx].float().mean((0, 2, 3)).view(1, -1, 1, 1) / 255 |
| if_std = IF[tr_idx].float().std((0, 2, 3)).view(1, -1, 1, 1) / 255 + 1e-3 |
| mk_mean, mk_std = MK[tr_idx].mean(0), MK[tr_idx].std(0) + 1e-3 |
|
|
|
|
| def batch(idx, d=None, train=False): |
| """Returns dict of modality tensors on GPU. d: dihedral id per sample (shared across image modalities).""" |
| idx = torch.as_tensor(idx) |
| if d is None: |
| d = torch.zeros(len(idx), dtype=torch.long) |
| x_if = IF[idx].to(dev, non_blocking=True).float() / 255 |
| x_nu = NUC[idx].to(dev, non_blocking=True).float() / 255 |
| x_if = torch.stack([dihedral(a, int(k)) for a, k in zip(x_if, d)]) |
| x_nu = torch.stack([dihedral(a, int(k)) for a, k in zip(x_nu, d)]) |
| if train: |
| x_if = x_if * (1 + 0.1 * torch.randn(len(idx), x_if.shape[1], 1, 1, device=dev)) |
| x_if = (x_if - if_mean.to(dev)) / if_std.to(dev) |
| x_nu = (x_nu - 0.5) / 0.5 |
| mk = ((MK[idx] - mk_mean) / mk_std).to(dev) |
| if train: |
| mk = mk + 0.05 * torch.randn_like(mk) |
| he = he_emb[idx, d].to(dev) |
| return {"he": he, "if": x_if, "nuclei": x_nu, "markers": mk} |
|
|
|
|
| |
| class ImgEnc(nn.Module): |
| def __init__(self, in_chans): |
| super().__init__() |
| self.trunk = timm.create_model(args.backbone, pretrained=True, in_chans=in_chans, num_classes=0) |
| self.head = nn.Sequential(nn.LayerNorm(self.trunk.num_features), nn.Linear(self.trunk.num_features, D)) |
|
|
| def forward(self, x): |
| return self.head(self.trunk(x)) |
|
|
|
|
| class MLPEnc(nn.Module): |
| def __init__(self, d_in): |
| super().__init__() |
| self.net = nn.Sequential(nn.Linear(d_in, 512), nn.GELU(), nn.Dropout(0.1), nn.Linear(512, 512), nn.GELU(), |
| nn.LayerNorm(512), nn.Linear(512, D)) |
|
|
| def forward(self, x): |
| return self.net(x) |
|
|
|
|
| class Bind(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.enc = nn.ModuleDict({"if": ImgEnc(17), "nuclei": ImgEnc(2), "markers": MLPEnc(MK.shape[1])}) |
| pairs = [("he", m) for m in ["if", "nuclei", "markers"]] |
| if args.loss == "allpairs": |
| pairs += [("if", "nuclei"), ("if", "markers"), ("nuclei", "markers")] |
| self.pairs = pairs |
| |
| self.logit_scale = nn.ParameterDict({f"{a}-{b}": nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32)) |
| for a, b in pairs}) |
|
|
| def embed(self, b): |
| z = {"he": F.normalize(b["he"], dim=-1)} |
| for m, e in self.enc.items(): |
| z[m] = F.normalize(e(b[m]), dim=-1) |
| return z |
|
|
| def loss(self, z): |
| tot, logs = 0, {} |
| tgt = torch.arange(len(z["he"]), device=dev) |
| for a, bm in self.pairs: |
| s = self.logit_scale[f"{a}-{bm}"].exp().clamp(max=100) |
| lg = s * z[a] @ z[bm].T |
| l = (F.cross_entropy(lg, tgt) + F.cross_entropy(lg.T, tgt)) / 2 |
| tot = tot + l |
| logs[f"{a}-{bm}"] = l.item() |
| return tot, logs |
|
|
|
|
| model = Bind().to(dev) |
| opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.05) |
| steps = args.epochs * max(1, len(tr_idx) // args.bs) |
| sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=steps, pct_start=0.05) |
| scaler = torch.amp.GradScaler() |
|
|
|
|
| @torch.no_grad() |
| def embed_all(idx): |
| model.eval() |
| zs = {m: [] for m in MODALITIES} |
| for s in range(0, len(idx), 256): |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| z = model.embed(batch(idx[s:s + 256])) |
| for m in MODALITIES: |
| zs[m].append(z[m].float().cpu()) |
| return {m: torch.cat(v) for m, v in zs.items()} |
|
|
|
|
| def retrieval_metrics(zq, zg): |
| sim = (zq @ zg.T).numpy() |
| n = len(sim) |
| rank = (sim > sim[np.arange(n), np.arange(n)][:, None]).sum(1) |
| r, c = linear_sum_assignment(-sim) |
| return {"R@1": float((rank < 1).mean()), "R@5": float((rank < 5).mean()), "R@10": float((rank < 10).mean()), |
| "MRR": float((1 / (rank + 1)).mean()), "MedR": float(np.median(rank + 1)), |
| "pair_acc": float((c[np.argsort(r)] == np.arange(n)).mean()), "N": n, "chance_R@1": 1 / n} |
|
|
|
|
| def fixed_gallery(zq, zg, size=100, reps=50): |
| """R@1 / pairing accuracy averaged over random galleries of a fixed size (comparable across splits).""" |
| rng = np.random.default_rng(0) |
| n = len(zq) |
| if n <= size: |
| m = retrieval_metrics(zq, zg) |
| return {"R@1_g100": m["R@1"], "pair_acc_g100": m["pair_acc"]} |
| r1, pa = [], [] |
| for _ in range(reps): |
| s = torch.from_numpy(rng.choice(n, size, replace=False)) |
| m = retrieval_metrics(zq[s], zg[s]) |
| r1.append(m["R@1"]); pa.append(m["pair_acc"]) |
| return {"R@1_g100": float(np.mean(r1)), "pair_acc_g100": float(np.mean(pa))} |
|
|
|
|
| def evaluate(idx, fixed=False): |
| z = embed_all(idx) |
| res = {} |
| for a, b in itertools.permutations(MODALITIES, 2): |
| res[f"{a}->{b}"] = retrieval_metrics(z[a], z[b]) |
| if fixed: |
| res[f"{a}->{b}"].update(fixed_gallery(z[a], z[b])) |
| return res, z |
|
|
|
|
| def mean_r1(res): |
| return float(np.mean([v["R@1"] for v in res.values()])) |
|
|
|
|
| |
| best, hist = -1, [] |
| t0 = time.time() |
| for ep in range(args.epochs): |
| model.train() |
| perm = np.random.permutation(tr_idx) |
| tl = [] |
| for s in range(0, len(perm) - args.bs + 1, args.bs): |
| idx = perm[s:s + args.bs] |
| b = batch(idx, d=torch.randint(0, 8, (len(idx),)), train=True) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| loss, logs = model.loss(model.embed(b)) |
| opt.zero_grad(set_to_none=True) |
| loss.backward() |
| opt.step(); sched.step() |
| tl.append(loss.item()) |
| vres, _ = evaluate(va_idx) |
| v = mean_r1(vres) |
| hist.append({"epoch": ep, "train_loss": float(np.mean(tl)), "val_meanR1": v, **{k: x for k, x in logs.items()}}) |
| print(f"ep {ep} loss {np.mean(tl):.3f} val meanR@1 {v:.3f} ({time.time() - t0:.0f}s)", flush=True) |
| if v > best: |
| best = v |
| torch.save(model.state_dict(), os.path.join(run_dir, "best.pt")) |
|
|
| |
| model.load_state_dict(torch.load(os.path.join(run_dir, "best.pt"))) |
| tres, z = evaluate(te_idx, fixed=True) |
| stats = pd.read_csv(os.path.join(os.path.dirname(RUNS), "results", "tile_stats.csv")).set_index("tile_id") |
| tissue = stats.loc[te_idx, "n_nuclei"].values >= args.min_nuclei_tissue |
| tres_tissue = {} |
| for a, b in itertools.permutations(MODALITIES, 2): |
| tres_tissue[f"{a}->{b}"] = retrieval_metrics(z[a][tissue], z[b][tissue]) |
| tres_tissue[f"{a}->{b}"].update(fixed_gallery(z[a][tissue], z[b][tissue])) |
| out = {"args": vars(args), "n_train": len(tr_idx), "n_val": len(va_idx), "n_test": len(te_idx), |
| "n_test_tissue": int(tissue.sum()), "best_val_meanR1": best, "test": tres, "test_tissue": tres_tissue, |
| "history": hist} |
| json.dump(out, open(os.path.join(run_dir, "metrics.json"), "w"), indent=1) |
| torch.save({m: v for m, v in z.items()}, os.path.join(run_dir, "test_embeddings.pt")) |
| print("TEST mean R@1 (all pairs):", mean_r1(tres), "| tissue-only:", mean_r1(tres_tissue)) |
| for k in ["he->if", "he->nuclei", "he->markers", "if->nuclei", "if->markers", "nuclei->markers"]: |
| print(f" {k:18s} R@1 {tres[k]['R@1']:.3f} R@5 {tres[k]['R@5']:.3f} pair_acc {tres[k]['pair_acc']:.3f} R@1_g100 {tres[k]['R@1_g100']:.3f}") |
|
|
