| """Frozen text-tower embeddings of the templated tile descriptions. |
| Usage: python 08_extract_text.py {plip|biomedclip|imagebind|conchclip} [--set <name> | --pathgen | --prompts] |
| Output: cache/text_<enc>.npy (N, 4 paraphrases, D) float16. |
| plip : vinid/plip CLIP text tower (77 tokens), projected -> 512 |
| biomedclip : microsoft/BiomedCLIP PubMedBERT text tower (256 tokens), projected -> 512 |
| imagebind : ImageBind-Huge text tower (OpenCLIP ViT-H text, 77 tokens) -> 1024 |
| conchclip : MahmoodLab/CONCH text tower (128 tokens), projected -> 512 |
| """ |
| import json |
| import os |
| import sys |
|
|
| import numpy as np |
| import torch |
|
|
| from common import CACHE, ROOT, 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 |
| PROMPTS = "--prompts" in sys.argv |
| if PROMPTS: |
| pr = json.load(open(os.path.join(ROOT, "prompts.json"))) |
| flat = [v["text"] for v in pr.values()] |
| ids, NV = list(range(len(flat))), 1 |
| elif PATHGEN: |
| import pandas as pd |
| flat = pd.read_csv(os.path.join(CACHE, f"{SET}_meta.csv")).caption.tolist() |
| ids, NV = list(range(len(flat))), 1 |
| else: |
| texts = json.load(open(os.path.join(CACHE, "text.json"))) |
| ids, NV = sorted(int(k) for k in texts), 4 |
| flat = [texts[str(i)][v] for i in ids for v in range(NV)] |
|
|
| if name == "plip": |
| from transformers import CLIPModel, CLIPTokenizer |
| m = CLIPModel.from_pretrained("vinid/plip").cuda().eval() |
| tok = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32") |
| enc = lambda b: m.get_text_features(**{k: v.cuda() for k, v in tok(b, padding=True, truncation=True, max_length=77, return_tensors="pt").items()}) |
| elif name == "biomedclip": |
| import open_clip |
| m, _ = open_clip.create_model_from_pretrained("hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224") |
| tok = open_clip.get_tokenizer("hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224") |
| m = m.cuda().eval() |
| enc = lambda b: m.encode_text(tok(b, context_length=256).cuda()) |
| elif name == "imagebind": |
| model, MT = load_imagebind_vision() |
| from imagebind.models.multimodal_preprocessors import SimpleTokenizer |
| tk = SimpleTokenizer(bpe_path=os.path.join(ROOT, "ImageBind_src", "imagebind", "bpe", "bpe_simple_vocab_16e6.txt.gz")) |
| enc = lambda b: model({MT.TEXT: torch.cat([tk(t).unsqueeze(0) for t in b]).cuda()})[MT.TEXT] |
| elif name == "conchclip": |
| from conch.open_clip_custom import create_model_from_pretrained, get_tokenizer, tokenize |
| m, _ = create_model_from_pretrained("conch_ViT-B-16", "hf_hub:MahmoodLab/conch") |
| m = m.cuda().eval() |
| tok = get_tokenizer() |
| enc = lambda b: m.encode_text(tokenize(tok, b).cuda(), normalize=False) |
| else: |
| raise ValueError(name) |
|
|
| out = [] |
| with torch.no_grad(): |
| for s in range(0, len(flat), 256): |
| out.append(enc(flat[s:s + 256]).float().cpu().numpy()) |
| out = np.concatenate(out).reshape(len(ids), NV, -1).astype(np.float16) |
| assert ids == list(range(len(ids))) |
| pre = "prompts_" if PROMPTS else f"{SET}_" if PATHGEN else "" |
| np.save(os.path.join(CACHE, pre + f"text_{name}.npy"), out) |
| print("saved", name, out.shape, np.isfinite(out).all()) |
|
|
