File size: 5,880 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 | """MIPHEI-ViT (Sanofi, arXiv 2505.10294) virtual mIF: H&E -> 16 ORION channels (all ORION markers except PD-1).
Model : ext/miphei/weights (GitHub release v1.0.0: decoder + LoRA, 26 MB; non-commercial academic licence) on top of
bioptimus/H-optimus-0 (HF, gated; HF token only via environment). model.py of the release (read in full) is
loaded from its path; nothing else from the downloaded repo is executed.
Training data of MIPHEI = this very ORION tile set with the ORION official split (37 train / 2 val / 2 test slides =
our splits/orion_official) -> MIPHEI virtual mIF is held-out ONLY on the official test (and val) tiles.
Input : RGB at 0.5 um/px (ORION tiles 333 px = 0.4997 um/px used as is; TCGA 672-px patches at ~0.25 um/px -> 336 px),
ImageNet mean/std (release config), reflect-padded to 512 px (input must be a power of 2), output cropped back.
Output: tanh in [-0.9, 0.9] -> (y + 0.9) / 1.8 * 255 = the 8-bit scale of the real ORION mIF tiles (MIPHEI targets).
cache/miphei/<set>_maps168.npy uint8 (N, 16, 168, 168) resampled exactly like the real tiles in 14_ (antialiased
bilinear to ~1 um/px) -> VirTues input
cache/miphei/<set>_means.npy float32 (N, 16) tile means at full resolution (0-255 scale)
cache/miphei/<set>_meta.csv, channels.json
Usage: HF_TOKEN=... .venv/bin/python 35_miphei_infer.py --set orion|tcga_text|dlpfc|idc [--bs 8]
"""
import argparse
import importlib.util
import json
import os
import time
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image, PngImagePlugin
from safetensors.torch import load_file
from common import CACHE, ORION, ROOT, load_all_tiles
p = argparse.ArgumentParser()
p.add_argument("--set", required=True)
p.add_argument("--bs", type=int, default=8)
args = p.parse_args()
dev = "cuda"
W = os.path.join(ROOT, "ext", "miphei", "weights")
OUT = os.path.join(CACHE, "miphei")
os.makedirs(OUT, exist_ok=True)
CH = ["Hoechst", "CD31", "CD45", "CD68", "CD4", "FOXP3", "CD8a", "CD45RO", "CD20", "PD-L1", "CD3e", "CD163",
"E-cadherin", "Ki67", "Pan-CK", "SMA"] # release config targ_channel_names (order of the output heads)
MEAN = torch.tensor([123.675, 116.28, 103.53], device=dev).view(1, 3, 1, 1) / 255
STD = torch.tensor([58.395, 57.12, 57.375], device=dev).view(1, 3, 1, 1) / 255
PAD_TO, VSIZE = 512, 168
PngImagePlugin.MAX_TEXT_CHUNK = 64 * 1024 ** 2 # some TCGA PNGs carry large text metadata chunks
spec = importlib.util.spec_from_file_location("miphei_model", os.path.join(W, "model.py"))
Mm = importlib.util.module_from_spec(spec); spec.loader.exec_module(Mm)
vit = Mm.get_hoptimus0_hf("bioptimus/H-optimus-0")
vit.set_input_size(img_size=(PAD_TO, PAD_TO))
enc = Mm.Encoder(vit)
model = Mm.MIPHEIViT(enc, Mm.Detail_Capture(emb_chans=enc.embed_dim, out_chans=len(CH), use_attention=True, activation=nn.Tanh()))
sd = Mm.merge_lora_weights(model, load_file(os.path.join(W, "model.safetensors")))
Mm.validate_load_info(model.load_state_dict(sd, strict=False))
model = model.eval().to(dev)
if args.set == "orion":
df = load_all_tiles()
paths, keys, scale = [os.path.join(ORION, x) for x in df.image_path], df.image_path.values, 1.0
elif args.set == "tcga_text":
I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
paths, keys, scale = list(I.image_path), I.key.values, 0.5
elif args.set in ("dlpfc", "idc"): # external sets (36_external_prep.py): uint8 RGB patches at 0.5 um/px, 224 px
E = os.path.join(CACHE, "external")
ARR = np.load(os.path.join(E, "dlpfc_he_05um.npy" if args.set == "dlpfc" else "idc_he.npy"), mmap_mode="r")
em = pd.read_csv(os.path.join(E, f"{args.set}_meta.csv"))
assert len(em) == len(ARR)
paths, keys, scale = list(range(len(ARR))), em.index.values, 1.0
else:
raise ValueError(args.set)
N = len(paths)
maps = np.lib.format.open_memmap(os.path.join(OUT, f"{args.set}_maps168.npy"), "w+", np.uint8, (N, len(CH), VSIZE, VSIZE))
means = np.zeros((N, len(CH)), np.float32)
sizes = []
t0 = time.time()
def load(pth):
if isinstance(pth, int):
return np.asarray(ARR[pth])
im = Image.open(pth).convert("RGB")
if scale != 1.0:
im = im.resize((round(im.width * scale), round(im.height * scale)), Image.BICUBIC)
return np.asarray(im)
with torch.inference_mode():
for s in range(0, N, args.bs):
ims = [load(x) for x in paths[s:s + args.bs]]
h, w = ims[0].shape[:2]
assert all(i.shape[:2] == (h, w) for i in ims) and h <= PAD_TO and w <= PAD_TO
x = torch.from_numpy(np.stack(ims)).to(dev).permute(0, 3, 1, 2).float() / 255
x = (x - MEAN) / STD
pt, pl = (PAD_TO - h) // 2, (PAD_TO - w) // 2
x = F.pad(x, (pl, PAD_TO - w - pl, pt, PAD_TO - h - pt), mode="reflect")
with torch.autocast("cuda", dtype=torch.float16): # as in the release notebook
y = model(x)
y = y.float().clamp(-0.9, 0.9)[:, :, pt:pt + h, pl:pl + w]
y = (y + 0.9) / 1.8 * 255
means[s:s + len(ims)] = y.mean((2, 3)).cpu().numpy()
maps[s:s + len(ims)] = F.interpolate(y, size=(VSIZE, VSIZE), mode="bilinear", antialias=True,
align_corners=False).round().clamp(0, 255).byte().cpu().numpy()
sizes += [(h, w)] * len(ims)
if s % (args.bs * 100) == 0:
print(f" {s}/{N} [{time.time() - t0:.0f}s]", flush=True)
maps.flush()
np.save(os.path.join(OUT, f"{args.set}_means.npy"), means)
pd.DataFrame({"key": keys, "path": paths, "h": [a for a, _ in sizes], "w": [b for _, b in sizes]}).to_csv(
os.path.join(OUT, f"{args.set}_meta.csv"), index=False)
json.dump({"channels": CH}, open(os.path.join(OUT, "channels.json"), "w"))
print(f"done {args.set}: {N} tiles [{time.time() - t0:.0f}s]")
|