| """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) |
|
|
