File size: 5,296 Bytes
e8fa066 | 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 | """Frozen pathology foundation-model embeddings of every H&E tile under the 8 dihedral transforms.
Usage: python 04b_extract_anchor_fm.py {uni2|conch|conchclip|gigapath|plip|phikonv2} [--set <name>]
Output: cache/he_<name>.npy (N, 8, D) float16 (raw features; L2-normalised later in training).
uni2 : MahmoodLab/UNI2-h, ViT-H/14 (reg tokens), 224px, ImageNet norm, CLS -> 1536-d
gigapath : prov-gigapath/prov-gigapath tile encoder, ViT-g/14, 224px, ImageNet norm -> 1536-d
conch : MahmoodLab/CONCH v1 vision tower (ViT-B/16), 448px, OpenAI-CLIP norm,
encode_image(proj_contrast=False, normalize=False) -> 512-d (authors' recommended tile feature)
plip : vinid/plip CLIP ViT-B/32 image features (projected) -> 512-d [open stand-in for CONCH]
phikonv2 : owkin/phikon-v2 ViT-L/16 DINOv2, CLS -> 1024-d [open stand-in for GigaPath]
"""
import os
import sys
import numpy as np
import timm
import torch
from PIL import Image, PngImagePlugin
PngImagePlugin.MAX_TEXT_CHUNK = 64 * 1024 * 1024 # some PathGen PNGs embed very large ICC profiles
from common import CACHE, ORION, dihedral, load_all_tiles, load_imagebind_vision
name = sys.argv[1]
SET = sys.argv[sys.argv.index("--set") + 1] if "--set" in sys.argv else "pathgen" if "--pathgen" in sys.argv else None
PATHGEN = SET is not None # encode an image-text pair set (cache/<SET>_meta.csv, e.g. PathGen-1.6M) instead of ORION tiles
UNI2_LOCAL = "/diskarray/home/tl688/reg2026/model_alldata/uni2-h.bin"
dev = "cuda"
IMNET = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
CLIPN = ([0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711])
if name == "uni2":
# HF repo is gated for this account -> use the local copy of the official UNI2-h checkpoint
model = timm.create_model(
"vit_giant_patch14_224", pretrained=False, img_size=224, patch_size=14, depth=24, num_heads=24,
init_values=1e-5, embed_dim=1536, mlp_ratio=2.66667 * 2, num_classes=0, no_embed_class=True,
mlp_layer=timm.layers.SwiGLUPacked, act_layer=torch.nn.SiLU, reg_tokens=8, dynamic_img_size=True)
model.load_state_dict(torch.load(UNI2_LOCAL, map_location="cpu"), strict=True)
size, norm, fwd = 224, IMNET, lambda m, x: m(x)
elif name == "gigapath":
model = timm.create_model("hf_hub:prov-gigapath/prov-gigapath", pretrained=True)
size, norm, fwd = 224, IMNET, lambda m, x: m(x)
elif name == "conch":
from conch.open_clip_custom import create_model_from_pretrained
model, _ = create_model_from_pretrained("conch_ViT-B-16", "hf_hub:MahmoodLab/conch")
size, norm = 448, CLIPN
fwd = lambda m, x: m.encode_image(x, proj_contrast=False, normalize=False)
elif name == "conchclip": # CONCH image tower projected into its joint image-text space (zero-shot with CONCH text)
from conch.open_clip_custom import create_model_from_pretrained
model, _ = create_model_from_pretrained("conch_ViT-B-16", "hf_hub:MahmoodLab/conch")
size, norm = 448, CLIPN
fwd = lambda m, x: m.encode_image(x, proj_contrast=True, normalize=False)
elif name == "plip": # open pathology CLIP (OpenPath tweets), stand-in for gated CONCH
from transformers import CLIPModel
model = CLIPModel.from_pretrained("vinid/plip")
size, norm, fwd = 224, CLIPN, lambda m, x: m.get_image_features(pixel_values=x)
elif name == "phikonv2": # open DINOv2 ViT-L pathology SSL model, stand-in for gated GigaPath
from transformers import AutoModel
model = AutoModel.from_pretrained("owkin/phikon-v2")
size, norm, fwd = 224, IMNET, lambda m, x: m(pixel_values=x).last_hidden_state[:, 0]
elif name == "biomedclip": # BiomedCLIP ViT-B/16 image tower (PMC-15M), projected -> 512; zero-shot baselines
import open_clip
model, _ = open_clip.create_model_from_pretrained("hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224")
size, norm, fwd = 224, CLIPN, lambda m, x: m.encode_image(x)
elif name == "imagebind": # used here for PathGen zero-shot baselines
model, MT = load_imagebind_vision()
size, norm, fwd = 224, CLIPN, lambda m, x: m({MT.VISION: x})[MT.VISION]
else:
raise ValueError(name)
model.eval().to(dev)
mean = torch.tensor(norm[0]).view(1, 3, 1, 1)
std = torch.tensor(norm[1]).view(1, 3, 1, 1)
if PATHGEN:
import pandas as pd
paths = pd.read_csv(os.path.join(CACHE, f"{SET}_meta.csv")).image_path.values
else:
paths = np.array([os.path.join(ORION, p) for p in load_all_tiles().image_path])
out = None
bs = 32
with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
for s in range(0, len(paths), bs):
ims = [np.asarray(Image.open(p).convert("RGB").resize((size, size), Image.BICUBIC))
for p in paths[s:s + bs]]
x = (torch.from_numpy(np.stack(ims)).permute(0, 3, 1, 2).float() / 255 - mean) / std
for d in range(8):
e = fwd(model, dihedral(x, d).to(dev)).float().cpu().numpy()
if out is None:
out = np.zeros((len(paths), 8, e.shape[1]), np.float16)
out[s:s + bs, d] = e
if s % 640 == 0:
print(name, s, len(paths), flush=True)
np.save(os.path.join(CACHE, (f"{SET}_" if PATHGEN else "") + f"he_{name}.npy"), out)
print("saved", name, out.shape, "finite:", np.isfinite(out).all())
|