File size: 10,372 Bytes
9a5439c | 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 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | """VirTues embeddings on the protein panel shared by ORION (real mIF) and GigaTIME (virtual mIF from H&E), so that one
mIF head trained on REAL ORION data can embed virtual mIF for validation (virtual data is never used for training).
Shared panel (ORION name <- GigaTIME channel): Hoechst<-DAPI, CD68, CD4, CD8a<-CD8, CD20, PD-L1, CD3e<-CD3, Ki67,
Pan-CK<-CK, SMA<-mean(Actin-D, Transgelin). PD-1 is left out (ORION PD-1 export is saturated / near constant).
Usage: python 30_virtues_shared.py real -> cache/if_virtues_shared.npy (3187, 8, 512) 8 dihedral views
python 30_virtues_shared.py virtual orion|tcga_text -> cache/if_virtues_shared_virtual_<set>.npy (N, 1, 512) view 0
python 30_virtues_shared.py virtual_qm orion,tcga_text -> cache/if_virtues_shared_virtual_qm_<set>_fold<k>.npy
(virtual pixels quantile-mapped per channel onto the real fold-k TRAIN pixel distribution)
--m16: 16-channel ORION panel shared with MIPHEI-ViT (all ORION markers but PD-1):
python 30_virtues_shared.py real --m16 -> cache/if_virtues_m16.npy (3187, 8, 512)
python 30_virtues_shared.py miphei <set> --m16 -> cache/if_virtues_m16_virtual_miphei_<set>.npy (N, 1, 512)
Preprocessing = 14_extract_if_fm.py (per-slide 0.99-quantile clip, log1p, dataset mean/std, 3x3 Gaussian blur, ~1 um/px):
real : ORION 333 px tiles (166 um) -> 168 px
virtual : GigaTIME maps (128 px; ORION tile 166 um, TCGA 672-px patch at ~0.25 um/px = 168 um) -> 168 px; values are
the uint8 sigmoid probabilities (0-255), standardised within their own set (slide = ORION slide / TCGA WSI)
"""
import importlib.util
import json
import os
import sys
import time
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from common import CACHE, ROOT, dihedral, load_all_tiles
MODE = sys.argv[1]
_argv, sys.argv = sys.argv, [sys.argv[0], "virtues", "--no_probe"] # 14_ parses its CLI at import time
spec = importlib.util.spec_from_file_location("virt14", os.path.join(ROOT, "14_extract_if_fm.py"))
V = importlib.util.module_from_spec(spec); spec.loader.exec_module(V)
sys.argv = _argv
dev, SIZE = V.dev, V.SIZE
SHARED = [("Hoechst", ["DAPI"]), ("CD68", ["CD68"]), ("CD4", ["CD4"]), ("CD8a", ["CD8"]), ("CD20", ["CD20"]),
("PD-L1", ["PD-L1"]), ("CD3e", ["CD3"]), ("Ki67", ["Ki67"]), ("Pan-CK", ["CK"]), ("SMA", ["Actin-D", "Transgelin"])]
ORION_NAMES = [m for m, _ in V.MARKERS]
UNIPROT = dict(V.MARKERS)
M16 = "--m16" in sys.argv # 16-channel panel shared by ORION and MIPHEI-ViT (all ORION markers except PD-1)
if M16:
sys.argv.remove("--m16")
SHARED = [(m, [m]) for m in ORION_NAMES if m != "PD-1"]
SUF = "m16" if M16 else "shared"
ids = [UNIPROT[m] for m, _ in SHARED]
t0 = time.time()
def load_virtual(st):
gch = json.load(open(os.path.join(CACHE, "gigatime", "channels.json")))["channels"]
maps = np.load(os.path.join(CACHE, "gigatime", f"{st}_maps.npy"), mmap_mode="r") # N,21,128,128 uint8
meta = pd.read_csv(os.path.join(CACHE, "gigatime", f"{st}_meta.csv"))
N = len(maps)
X = np.empty((N, len(SHARED), SIZE, SIZE), np.float16)
for s in range(0, N, 256):
m = torch.from_numpy(np.asarray(maps[s:s + 256])).to(dev).float()
x = torch.stack([m[:, [gch.index(g) for g in gs]].mean(1) for _, gs in SHARED], 1)
X[s:s + 256] = F.interpolate(x, size=(SIZE, SIZE), mode="bilinear", align_corners=False).clamp(min=0).half().cpu().numpy()
if st == "orion":
df = load_all_tiles()
assert (meta.tile_id.values == df.tile_id.values).all()
slide = df.in_slide_name.values
else:
I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
assert (meta.key.values == I.key.values).all()
slide = I.wsi_id.values
return X, slide
def encode(X, slide, views, out_f):
X, uq, mean, std, valid, slide_k, n_zero = V.standardize(X, slide)
print(" log1p mean", np.round(mean, 2).tolist(), "\n std", np.round(std, 2).tolist(), "| zero-uq (slide,ch):", int(n_zero.sum()))
model, midx = V.load_virtues(ids)
midx = midx.to(dev)
N = len(X)
out = np.zeros((N, views, 512), np.float16)
BS = 4 if views == 8 else 32 # 32 crops per forward (~4.7 GiB peak, shared GPU)
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
for s in range(0, N, BS):
x = torch.from_numpy(X[s:s + BS]).to(dev).float()
crops = [dihedral(x, d)[i] for d in range(views) for i in range(len(x))] # d-major
ps = model.encoder.forward_list(crops, [midx] * len(crops)).patch_summary_tokens
e = torch.stack([p.float().mean((0, 1)) for p in ps]).view(views, len(x), -1).transpose(0, 1)
out[s:s + BS] = e.cpu().numpy().astype(np.float16)
np.save(out_f, out)
print(f"saved {out_f} {out.shape} finite {np.isfinite(out).all()} [{time.time() - t0:.0f}s]", flush=True)
if MODE == "miphei":
# MIPHEI virtual mIF (35_miphei_infer.py): already on the 8-bit ORION scale and resampled to SIZE like the real tiles
assert M16
st = sys.argv[2]
ch = json.load(open(os.path.join(CACHE, "miphei", "channels.json")))["channels"]
X = np.load(os.path.join(CACHE, "miphei", f"{st}_maps168.npy"))[:, [ch.index(m) for m, _ in SHARED]].astype(np.float16)
if st == "orion":
slide = load_all_tiles().in_slide_name.values
else:
slide = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv")).wsi_id.values if st == "tcga_text" else \
pd.read_csv(os.path.join(CACHE, "external", f"{st}_meta.csv"))["sample"].astype(str).values
assert len(slide) == len(X)
encode(X, slide, 1, os.path.join(CACHE, f"if_virtues_m16_virtual_miphei_{st}.npy"))
sys.exit(0)
if MODE == "virtual_qm":
# channel-UNIFIED virtual mIF: per channel, quantile-map the GigaTIME probability pixels onto the REAL ORION pixel
# intensity distribution of the fold-k TRAINING tiles (monotone, keeps the spatial pattern), then the real
# preprocessing + VirTues. Only real training-split pixels define the target distribution; nothing is trained.
sets = sys.argv[2].split(",")
df = load_all_tiles()
X17, _ = V.load_resampled(df)
XR = X17[:, [ORION_NAMES.index(m) for m, _ in SHARED]]; del X17
P = np.linspace(0, 1, 1001)
rng = np.random.default_rng(0)
ref = {}
for k in range(5):
sp = pd.read_csv(os.path.join(ROOT, "splits", "patient_5fold", f"fold{k}.csv"))
tr = sp.tile_id[sp.split == "train"].values
sub = rng.choice(tr, min(400, len(tr)), replace=False)
ref[k] = np.stack([np.quantile(XR[sub, c].astype(np.float32).ravel(), P) for c in range(len(SHARED))])
del XR
np.savez(os.path.join(CACHE, "gigatime", "qm_ref_real_train.npz"), P=P, **{f"fold{k}": v for k, v in ref.items()})
for st in sets:
X0, slide = load_virtual(st)
sub = rng.choice(len(X0), min(400, len(X0)), replace=False)
vq = np.stack([np.quantile(X0[sub, c].astype(np.float32).ravel(), P) for c in range(len(SHARED))])
vq = vq + np.arange(len(P))[None] * 1e-6 # strictly increasing despite the many zero-probability ties
for k in range(5):
X = np.empty_like(X0)
for s in range(0, len(X0), 128):
x = X0[s:s + 128].astype(np.float32)
for c in range(len(SHARED)):
x[:, c] = np.interp(x[:, c], vq[c], ref[k][c])
X[s:s + 128] = x.astype(np.float16)
print(f"[{time.time() - t0:.0f}s] {st} fold {k}: quantile-mapped", flush=True)
encode(X, slide, 1, os.path.join(CACHE, f"if_virtues_shared_virtual_qm_{st}_fold{k}.npy"))
del X
sys.exit(0)
if MODE == "real":
df = load_all_tiles()
X17, _ = V.load_resampled(df)
X = np.ascontiguousarray(X17[:, [ORION_NAMES.index(m) for m, _ in SHARED]]); del X17
slide, views, out_f = df.in_slide_name.values, 8, os.path.join(CACHE, f"if_virtues_{SUF}.npy")
elif MODE == "virtual":
st = sys.argv[2]
gch = json.load(open(os.path.join(CACHE, "gigatime", "channels.json")))["channels"]
maps = np.load(os.path.join(CACHE, "gigatime", f"{st}_maps.npy"), mmap_mode="r") # N,21,128,128 uint8
meta = pd.read_csv(os.path.join(CACHE, "gigatime", f"{st}_meta.csv"))
N = len(maps)
X = np.empty((N, len(SHARED), SIZE, SIZE), np.float16)
for s in range(0, N, 256):
m = torch.from_numpy(np.asarray(maps[s:s + 256])).to(dev).float()
x = torch.stack([m[:, [gch.index(g) for g in gs]].mean(1) for _, gs in SHARED], 1)
X[s:s + 256] = F.interpolate(x, size=(SIZE, SIZE), mode="bilinear", align_corners=False).clamp(min=0).half().cpu().numpy()
if st == "orion":
df = load_all_tiles()
assert (meta.tile_id.values == df.tile_id.values).all()
slide = df.in_slide_name.values
else:
I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
assert (meta.key.values == I.key.values).all()
slide = I.wsi_id.values
views, out_f = 1, os.path.join(CACHE, f"if_virtues_shared_virtual_{st}.npy")
else:
raise ValueError(MODE)
print(f"[{time.time() - t0:.0f}s] loaded {X.shape}; standardising", flush=True)
X, uq, mean, std, valid, slide_k, n_zero = V.standardize(X, slide)
print(" log1p mean", np.round(mean, 2).tolist(), "\n std", np.round(std, 2).tolist(), "| zero-uq (slide,ch):", int(n_zero.sum()))
model, midx = V.load_virtues(ids)
midx = midx.to(dev)
N = len(X)
out = np.zeros((N, views, 512), np.float16)
BS = 4 if views == 8 else 32 # 32 crops per forward (~4.7 GiB peak, shared GPU)
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
for s in range(0, N, BS):
x = torch.from_numpy(X[s:s + BS]).to(dev).float()
crops = [dihedral(x, d)[i] for d in range(views) for i in range(len(x))] # d-major
ps = model.encoder.forward_list(crops, [midx] * len(crops)).patch_summary_tokens
e = torch.stack([p.float().mean((0, 1)) for p in ps]).view(views, len(x), -1).transpose(0, 1)
out[s:s + BS] = e.cpu().numpy().astype(np.float16)
if s % (BS * 50) == 0:
print(f" encoded {s}/{N} [{time.time() - t0:.0f}s]", flush=True)
np.save(out_f, out)
print(f"saved {out_f} {out.shape} finite {np.isfinite(out).all()} [{time.time() - t0:.0f}s]")
|