File size: 4,177 Bytes
c2ede1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | """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] # CRC33_01 / CRC33_02 -> CRC33
df["if_batch"] = df.targ_slide_path.str.extract(r"_(C59[a-z]X)_")[0] # IF acquisition batch
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 # other candidates tied with the true match
def rk(k): # P(true match within top-k) with random tie-breaking
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
|