File size: 1,068 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 | """Frozen ImageBind-Huge vision embeddings of every H&E tile under each of the 8 dihedral transforms.
Output cache/he_imagebind.npy : (N, 8, 1024) float16, L2-normalised (ImageBind postprocessor)."""
import os
import numpy as np
import torch
from common import CACHE, dihedral, load_imagebind_vision
MEAN = torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1)
STD = torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1)
he = np.load(os.path.join(CACHE, "he.npy"), mmap_mode="r")
model, MT = load_imagebind_vision()
out = np.zeros((len(he), 8, 1024), np.float16)
bs = 64
with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
for s in range(0, len(he), bs):
x = (torch.from_numpy(np.ascontiguousarray(he[s:s + bs])).float() / 255 - MEAN) / STD
for d in range(8):
e = model({MT.VISION: dihedral(x, d).cuda()})[MT.VISION]
out[s:s + bs, d] = e.float().cpu().numpy()
print(s, len(he), flush=True)
np.save(os.path.join(CACHE, "he_imagebind.npy"), out)
print("saved", out.shape)
|