File size: 2,967 Bytes
ef01023 | 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 | r'''
OpenVision2 ViT-H/14 @224 — image -> caption demo.
Combines the released vision encoder (this repo, open_clip format) with the caption
text decoder (this repo). Requires the patched open_clip that ships
`create_vision_encoder_and_transforms` (from the OpenVision repo:
https://github.com/UCSC-VLAA/OpenVision), plus torch / safetensors / pillow.
python caption_example.py --image path/to/img.jpg
'''
import argparse, json
import numpy as np, torch
from PIL import Image
from huggingface_hub import hf_hub_download
from open_clip.factory import create_vision_encoder_and_transforms # patched open_clip
from modeling_openvision2_decoder import OpenVision2TextDecoder, OpenVision2TextDecoderConfig
from safetensors.torch import load_file
REPO = "UCSC-VLAA/openvision2-vit-large-patch14-336-vision-only"
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406]) * 255
IMAGENET_STD = np.array([0.229, 0.224, 0.225]) * 255
def preprocess(path, res=224):
im = Image.open(path).convert("RGB")
w, h = im.size; s = res / min(w, h)
im = im.resize((round(w * s), round(h * s)), Image.BILINEAR)
w, h = im.size; l, t = (w - res) // 2, (h - res) // 2
im = im.crop((l, t, l + res, t + res))
x = (np.asarray(im, np.float32) - IMAGENET_MEAN) / IMAGENET_STD
return torch.tensor(x.transpose(2, 0, 1)[None], dtype=torch.float32) # [1,C,H,W]
def load_decoder(cache_dir=None):
cfg = json.load(open(hf_hub_download(REPO, "text_decoder_config.json", cache_dir=cache_dir)))
dec = OpenVision2TextDecoder(OpenVision2TextDecoderConfig(
width=cfg["width"], depth=cfg["depth"], num_heads=cfg["num_heads"], mlp_dim=cfg["mlp_dim"],
vocab_size=cfg["vocab_size"], vision_width=cfg["vision_width"]))
dec.load_state_dict(load_file(hf_hub_download(REPO, "caption_decoder.safetensors", cache_dir=cache_dir)))
dec.eval()
vocab = [l.rstrip("\n") for l in open(hf_hub_download(REPO, "bert_base_vocab_bos_eos.txt", cache_dir=cache_dir))]
return dec, vocab
def detok(ids, vocab):
w = []
for i in ids:
if i == 2: break # eos
if i in (0, 1, 2): continue # pad / bos / eos
tk = vocab[i] if 0 <= i < len(vocab) else "[UNK]"
if tk.startswith("##"): w[-1] = (w[-1] + tk[2:]) if w else tk[2:]
else: w.append(tk)
return " ".join(w)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--image", required=True)
ap.add_argument("--max_len", type=int, default=64)
ap.add_argument("--cache_dir", default=None)
args = ap.parse_args()
enc = create_vision_encoder_and_transforms(model_name=f"hf-hub:{REPO}", cache_dir=args.cache_dir).eval()
dec, vocab = load_decoder(args.cache_dir)
with torch.no_grad():
_, tokens = enc(preprocess(args.image)) # [1, N, 1280]
ids = dec.generate(tokens, max_len=args.max_len, bos_id=1, eos_id=2)[0].tolist()
print(detok(ids, vocab))
if __name__ == "__main__":
main()
|