| """Shared paths / helpers for the ORION ImageBind-style 4-modality experiments. |
| |
| Modalities (all spatially paired per 333x333 tile at 20x): |
| he : H&E RGB tile -> anchor, frozen ImageBind-Huge vision encoder |
| if : 17-channel multiplex IF (uint8) -> trainable encoder |
| nuclei : nuclei instance mask (int32) -> trainable encoder (binary mask + boundary channels) |
| markers : per-tile marker vector (17 *_prop + 16 *_count from the dataframes) -> trainable MLP |
| """ |
| import os |
| import sys |
| import types |
|
|
| import numpy as np |
| import pandas as pd |
|
|
| ROOT = os.path.dirname(os.path.abspath(__file__)) |
| ORION = os.path.join(ROOT, "..", "ORION", "ORIONCRC_dataset_tile_20x") |
| CACHE = os.path.join(ROOT, "cache") |
| SPLITS = os.path.join(ROOT, "splits") |
| RUNS = os.path.join(ROOT, "runs") |
| IB_WEIGHTS = os.path.join(ROOT, "weights", "imagebind_huge.pth") |
| IMG = 224 |
|
|
| MODALITIES = ["he", "if", "nuclei", "markers"] |
|
|
|
|
| def load_all_tiles(): |
| """All tiles across the official train/val/test csvs, with a stable global index.""" |
| dfs = [] |
| for s in ["train", "val", "test"]: |
| d = pd.read_csv(os.path.join(ORION, f"{s}_dataframe.csv")) |
| d["orion_split"] = s |
| dfs.append(d) |
| df = pd.concat(dfs, ignore_index=True) |
| slides = pd.read_csv(os.path.join(ORION, "slide_dataframe.csv")) |
| df = df.merge(slides[["in_slide_name", "orion_slide_id", "targ_slide_path"]], on="in_slide_name", how="left") |
| df["patient"] = df.orion_slide_id.str.split("_").str[0] |
| df["if_batch"] = df.targ_slide_path.str.extract(r"_(C59[a-z]X)_")[0] |
| df["tile_id"] = np.arange(len(df)) |
| return df |
|
|
|
|
| def marker_columns(df): |
| prop = [c for c in df.columns if c.endswith("_prop")] |
| cnt = [c for c in df.columns if c.endswith("_count")] |
| return prop, cnt |
|
|
|
|
| def marker_matrix(df): |
| prop, cnt = marker_columns(df) |
| return np.concatenate([df[prop].values, np.log1p(df[cnt].values)], 1).astype(np.float32) |
|
|
|
|
| def dihedral(x, d): |
| """Apply one of 8 dihedral transforms to the last two (spatial) dims of a tensor/array.""" |
| import torch |
| k, flip = d % 4, d // 4 |
| if flip: |
| x = torch.flip(x, dims=(-1,)) |
| return torch.rot90(x, k, dims=(-2, -1)) |
|
|
|
|
| def load_imagebind_vision(device="cuda"): |
| """Build ImageBind-Huge and load the official checkpoint without importing imagebind.data |
| (which needs torchaudio/pytorchvideo).""" |
| src = os.path.join(ROOT, "ImageBind_src", "imagebind") |
| if "imagebind" not in sys.modules: |
| pkg = types.ModuleType("imagebind") |
| pkg.__path__ = [src] |
| sys.modules["imagebind"] = pkg |
| import torch |
| from imagebind.models import imagebind_model |
| from imagebind.models.imagebind_model import ModalityType |
|
|
| model = imagebind_model.imagebind_huge(pretrained=False) |
| sd = torch.load(IB_WEIGHTS, map_location="cpu") |
| missing, unexpected = model.load_state_dict(sd, strict=True), None |
| model.eval().to(device) |
| return model, ModalityType |
|
|
|
|
| def metrics(zq, zg, g100=True, eps=1e-6): |
| """Tie-aware retrieval metrics (expected rank under random tie-breaking; template texts can tie).""" |
| import torch |
| from scipy.optimize import linear_sum_assignment |
| sim = (zq @ zg.T).numpy() |
| n = len(sim) |
| d = sim[np.arange(n), np.arange(n)][:, None] |
| gt = (sim > d + eps).sum(1) |
| tie = (np.abs(sim - d) <= eps).sum(1) - 1 |
|
|
| def rk(k): |
| return np.clip((k - gt) / (tie + 1), 0, 1).mean() |
|
|
| r, c = linear_sum_assignment(-sim) |
| res = {"R@1": rk(1), "R@5": rk(5), "R@10": rk(10), "MRR": float(np.mean(1 / (gt + tie / 2 + 1))), |
| "pair_acc": float((c[np.argsort(r)] == np.arange(n)).mean()), "N": n} |
| if g100: |
| rng = np.random.default_rng(0) |
| if n <= 100: |
| res["R@1_g100"] = res["R@1"] |
| else: |
| res["R@1_g100"] = float(np.mean([metrics(zq[s], zg[s], g100=False)["R@1"] |
| for s in (torch.from_numpy(rng.choice(n, 100, replace=False)) for _ in range(50))])) |
| return res |
|
|
