| """MedVTok (general medical visual tokenizer, MERA-aligned to CONCH v1.5 for histopathology) global image features |
| of every H&E set, as an additional / alternative image encoder next to the frozen UNI2 anchor. |
| |
| Run: .venv/bin/python -I 50_medvtok_feats.py <set> [<set> ...] sets: orion tcga hest dlpfc |
| Model: ext/medvtok/repo (Masaaki-75/medvtok @a558c24, MIT; only models/medvtok.py + models/utils.py + |
| models/transformer.py are imported) + ext/medvtok/weights/ckpt.pth (HF massaki75/medvtok, MIT), loaded here with |
| torch.load(weights_only=True) and strict=True (the repo's build_model uses weights_only=False / strict=False). |
| Encoder only (decode=False). Preprocessing = repo build_transform: RGB -> bicubic resize 256 -> [-1, 1]. |
| Feature = outputs["feats"]: mean-pooled encoder tokens -> LayerNorm -> Linear (1152-d, the aligned global feature). |
| Views: orion / tcga 8 dihedral transforms (= he_uni2.npy); hest / dlpfc 1 view (= their UNI2 caches). |
| Images: orion ../ORION tiles, tcga cache/allimg_meta.csv image_path, hest cache/hest/he_img.npy (49_), |
| dlpfc cache/external/dlpfc_he_05um.npy (0.5 um/px, as for MIPHEI / captions). |
| Output: cache/he_medvtok.npy (3187, 8, 1152), cache/allimg_he_medvtok.npy (6144, 8, 1152), |
| cache/hest/he_medvtok.npy (9484, 1, 1152), cache/external/dlpfc_he_medvtok.npy (482, 1152); float16. |
| """ |
| import json |
| import os |
| import sys |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import torch.nn.functional as F |
| from PIL import Image, PngImagePlugin |
|
|
| PngImagePlugin.MAX_TEXT_CHUNK = 64 * 1024 * 1024 |
|
|
| ROOT = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, ROOT) |
| from common import CACHE, ORION, dihedral, load_all_tiles |
|
|
| REPO = os.path.join(ROOT, "ext", "medvtok", "repo") |
| WTS = os.path.join(ROOT, "ext", "medvtok", "weights") |
| sys.path.insert(0, REPO) |
| from models.medvtok import MedVTok |
|
|
| dev = "cuda" |
| SIZE, BS = 256, 64 |
|
|
| cfg = json.load(open(os.path.join(WTS, "config.json"))) |
| assert cfg["architecture"] == "medvtok-clip-meanpool" |
| model = MedVTok(**cfg["model_kwargs"]) |
| sd = torch.load(os.path.join(WTS, "ckpt.pth"), map_location="cpu", weights_only=True) |
| print(model.load_state_dict(sd, strict=True)) |
| del sd, model.decoder |
| model.eval().to(dev) |
|
|
|
|
| @torch.no_grad() |
| def encode(u8): |
| """u8: (B, H, W, 3) uint8 -> (B, 1152) float32""" |
| x = torch.from_numpy(np.ascontiguousarray(u8)).to(dev).permute(0, 3, 1, 2).float() / 255 |
| x = F.interpolate(x, size=(SIZE, SIZE), mode="bicubic", align_corners=False, antialias=True).clamp(0, 1) |
| x = x * 2 - 1 |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| out = model([im for im in x], sample_posterior=False, decode=False, return_batched_tensor=False) |
| return out["feats"].float().cpu().numpy() |
|
|
|
|
| def run(load, n, views): |
| out = np.zeros((n, views, 1152), np.float16) |
| for s in range(0, n, BS): |
| u8 = load(s, min(n, s + BS)) |
| for d in range(views): |
| x = u8 if views == 1 else dihedral(torch.from_numpy(u8).permute(0, 3, 1, 2), d).permute(0, 2, 3, 1).numpy() |
| out[s:s + len(u8), d] = encode(x) |
| if s % (BS * 20) == 0: |
| print(f" {s}/{n}", flush=True) |
| assert np.isfinite(out.astype(np.float32)).all() |
| return out |
|
|
|
|
| def from_paths(paths): |
| return lambda a, b: np.stack([np.asarray(Image.open(p).convert("RGB").resize((SIZE, SIZE), Image.BICUBIC)) |
| for p in paths[a:b]]) |
|
|
|
|
| for st in sys.argv[1:]: |
| if st == "orion": |
| paths = np.array([os.path.join(ORION, p) for p in load_all_tiles().image_path]) |
| np.save(os.path.join(CACHE, "he_medvtok.npy"), run(from_paths(paths), len(paths), 8)) |
| elif st == "tcga": |
| paths = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv")).image_path.values |
| np.save(os.path.join(CACHE, "allimg_he_medvtok.npy"), run(from_paths(paths), len(paths), 8)) |
| elif st == "hest": |
| A = np.load(os.path.join(CACHE, "hest", "he_img.npy"), mmap_mode="r") |
| np.save(os.path.join(CACHE, "hest", "he_medvtok.npy"), run(lambda a, b: np.asarray(A[a:b]), len(A), 1)) |
| elif st == "dlpfc": |
| A = np.load(os.path.join(CACHE, "external", "dlpfc_he_05um.npy")) |
| np.save(os.path.join(CACHE, "external", "dlpfc_he_medvtok.npy"), run(lambda a, b: A[a:b], len(A), 1)[:, 0]) |
| print("saved", st, flush=True) |
|
|
