|
|
--- |
|
|
license: apache-2.0 |
|
|
--- |
|
|
|
|
|
# OFA-base-caption |
|
|
This is the **base** version of OFA model finetuned for the image captioning task. OFA is a unified multimodal pretrained model that unifies modalities (i.e., cross-modality, vision, language) and tasks (e.g., image generation, visual grounding, image captioning, image classification, text generation, etc.) to a simple sequence-to-sequence learning framework. |
|
|
|
|
|
The directory includes 4 files, namely `config.json` which consists of model configuration, `vocab.json` and `merge.txt` for our OFA tokenizer, and lastly `pytorch_model.bin` which consists of model weights. There is no need to worry about the mismatch between Fairseq and transformers, since we have addressed the issue yet. |
|
|
|
|
|
To use it in transformers, please refer to https://github.com/OFA-Sys/OFA/tree/feature/add_transformers. Install the transformers and download the models as shown below. |
|
|
|
|
|
``` |
|
|
git clone --single-branch --branch feature/add_transformers https://github.com/OFA-Sys/OFA.git |
|
|
pip install OFA/transformers/ |
|
|
``` |
|
|
|
|
|
After, prepare an image for the testing example below. Also, ensure that you have pillow and torchvision in your environment. |
|
|
|
|
|
``` |
|
|
import re |
|
|
import time |
|
|
from PIL import Image |
|
|
from torchvision import transforms |
|
|
from transformers import OFATokenizer, OFAModel |
|
|
|
|
|
model_name = "OFA-sys/OFA-base-caption" |
|
|
|
|
|
mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5] |
|
|
resolution = 256 |
|
|
|
|
|
patch_resize_transform = transforms.Compose([ |
|
|
lambda image: image.convert("RGB"), |
|
|
transforms.Resize((resolution, resolution), interpolation=Image.BICUBIC), |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize(mean=mean, std=std) |
|
|
]) |
|
|
|
|
|
start = time.time() |
|
|
tokenizer = OFATokenizer.from_pretrained(model_name) |
|
|
model = OFAModel.from_pretrained(model_name, use_cache=False) |
|
|
alapsed = time.time() - start |
|
|
print(f"Loaded in {alapsed} secs") |
|
|
|
|
|
|
|
|
def caption_image(txt, img): |
|
|
inputs = tokenizer([txt], return_tensors="pt").input_ids |
|
|
patch_img = patch_resize_transform(img).unsqueeze(0) |
|
|
|
|
|
gen = model.generate(inputs, patch_images=patch_img, num_beams=5, no_repeat_ngram_size=3) |
|
|
results = tokenizer.batch_decode(gen, skip_special_tokens=True) |
|
|
|
|
|
result = results[0].strip() |
|
|
result = re.sub(r'[^\w\s]', '', result) |
|
|
|
|
|
return result |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
txt = "What does the image describe?" |
|
|
img = Image.open('/path/to/input/image.jpg') |
|
|
caption = caption_image(txt, img) |
|
|
print(caption) |
|
|
|
|
|
``` |
|
|
|
