import os os.system('cd fairseq;' 'pip install ./; cd ..') os.system('ls -l') import torch import numpy as np import gradio as gr import cv2 from PIL import Image from torchvision import transforms from fairseq import utils, tasks, options from fairseq import checkpoint_utils from fairseq.dataclass.utils import convert_namespace_to_omegaconf from tasks.mm_tasks.caption import CaptionTask from tasks.mm_tasks.refcoco import RefcocoTask from tasks.mm_tasks.vqa_gen import VqaGenTask def move2gpu(models, cfg): for model in models: model.eval() if use_fp16: model.half() if use_cuda and not cfg.distributed_training.pipeline_model_parallel: model.cuda() model.prepare_for_inference_(cfg) def construct_transform(patch_image_size): mean = [0.5, 0.5, 0.5] std = [0.5, 0.5, 0.5] patch_resize_transform = transforms.Compose([ lambda image: image.convert("RGB"), transforms.Resize((patch_image_size, patch_image_size), interpolation=Image.BICUBIC), transforms.ToTensor(), transforms.Normalize(mean=mean, std=std), ]) return patch_resize_transform # Register tasks tasks.register_task('caption', CaptionTask) tasks.register_task('refcoco', RefcocoTask) tasks.register_task('vqa_gen', VqaGenTask) # turn on cuda if GPU is available use_cuda = torch.cuda.is_available() # use fp16 only when GPU is available use_fp16 = True # download checkpoints os.system('wget https://ofa-silicon.oss-us-west-1.aliyuncs.com/checkpoints/caption_demo.pt; ' 'mkdir -p checkpoints; mv caption_demo.pt checkpoints/caption_demo.pt') os.system('wget https://ofa-silicon.oss-us-west-1.aliyuncs.com/checkpoints/refcoco_demo.pt; ' 'mkdir -p checkpoints; mv refcoco_demo.pt checkpoints/refcoco_demo.pt') os.system('wget https://ofa-silicon.oss-us-west-1.aliyuncs.com/checkpoints/general_demo.pt; ' 'mkdir -p checkpoints; mv general_demo.pt checkpoints/general_demo.pt') # Load ckpt & config for Image Captioning caption_overrides = {"bpe_dir": "utils/BPE", "eval_cider": False, "beam": 5, "max_len_b": 16, "no_repeat_ngram_size": 3, "seed": 7} caption_models, caption_cfg, caption_task = checkpoint_utils.load_model_ensemble_and_task( utils.split_paths('checkpoints/caption_demo.pt'), arg_overrides=caption_overrides ) # Load ckpt & config for Refcoco refcoco_overrides = {"bpe_dir": "utils/BPE", "eval_cider": False, "beam": 5, "max_len_b": 16, "no_repeat_ngram_size": 3, "seed": 7} refcoco_models, refcoco_cfg, refcoco_task = checkpoint_utils.load_model_ensemble_and_task( utils.split_paths('checkpoints/refcoco_demo.pt'), arg_overrides=refcoco_overrides ) refcoco_cfg.common.seed = 7 refcoco_cfg.generation.beam = 5 refcoco_cfg.generation.min_len = 4 refcoco_cfg.generation.max_len_a = 0 refcoco_cfg.generation.max_len_b = 4 refcoco_cfg.generation.no_repeat_ngram_size = 3 # Load pretrained ckpt & config for VQA parser = options.get_generation_parser() input_args = ["", "--task=vqa_gen", "--beam=100", "--unnormalized", "--path=checkpoints/general_demo.pt", "--bpe-dir=utils/BPE"] args = options.parse_args_and_arch(parser, input_args) vqa_cfg = convert_namespace_to_omegaconf(args) vqa_task = tasks.setup_task(vqa_cfg.task) vqa_models, vqa_cfg = checkpoint_utils.load_model_ensemble( utils.split_paths(vqa_cfg.common_eval.path), task=vqa_task ) # Load pretrained ckpt & config for Generic Interface parser = options.get_generation_parser() input_args = ["", "--task=refcoco", "--beam=10", "--path=checkpoints/general_demo.pt", "--bpe-dir=utils/BPE", "--no-repeat-ngram-size=3", "--patch-image-size=384"] args = options.parse_args_and_arch(parser, input_args) general_cfg = convert_namespace_to_omegaconf(args) general_task = tasks.setup_task(general_cfg.task) general_models, general_cfg = checkpoint_utils.load_model_ensemble( utils.split_paths(general_cfg.common_eval.path), task=general_task ) # move models to gpu move2gpu(caption_models, caption_cfg) move2gpu(refcoco_models, refcoco_cfg) move2gpu(vqa_models, vqa_cfg) move2gpu(general_models, general_cfg) # Initialize generator caption_generator = caption_task.build_generator(caption_models, caption_cfg.generation) refcoco_generator = refcoco_task.build_generator(refcoco_models, refcoco_cfg.generation) vqa_generator = vqa_task.build_generator(vqa_models, vqa_cfg.generation) vqa_generator.zero_shot = True vqa_generator.constraint_trie = None general_generator = general_task.build_generator(general_models, general_cfg.generation) # Construct image transforms caption_transform = construct_transform(caption_cfg.task.patch_image_size) refcoco_transform = construct_transform(refcoco_cfg.task.patch_image_size) vqa_transform = construct_transform(vqa_cfg.task.patch_image_size) general_transform = construct_transform(general_cfg.task.patch_image_size) # Text preprocess bos_item = torch.LongTensor([caption_task.src_dict.bos()]) eos_item = torch.LongTensor([caption_task.src_dict.eos()]) pad_idx = caption_task.src_dict.pad() def get_symbols_to_strip_from_output(generator): if hasattr(generator, "symbols_to_strip_from_output"): return generator.symbols_to_strip_from_output else: return {generator.bos, generator.eos} def decode_fn(x, tgt_dict, bpe, generator, tokenizer=None): x = tgt_dict.string(x.int().cpu(), extra_symbols_to_ignore=get_symbols_to_strip_from_output(generator)) token_result = [] bin_result = [] img_result = [] for token in x.strip().split(): if token.startswith('