fourmodal-pathology-bind / experiments /30_virtues_shared.py
Carpediem123's picture
Migrate release files 181-210
9a5439c verified
Raw History Blame Contribute Delete
10.4 kB
"""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]")