caption_example.py

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-huge-patch14-448-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()