ICML2022/OFA

1import os2 3os.system('cd fairseq;'4          'pip install --use-feature=in-tree-build ./; cd ..')5os.system('ls -l')6 7import torch8import numpy as np9import gradio as gr10import cv211from PIL import Image12from torchvision import transforms13 14from fairseq import utils, tasks, options15from fairseq import checkpoint_utils16from fairseq.dataclass.utils import convert_namespace_to_omegaconf17 18from tasks.mm_tasks.caption import CaptionTask19from tasks.mm_tasks.refcoco import RefcocoTask20from tasks.mm_tasks.vqa_gen import VqaGenTask21 22 23def move2gpu(models, cfg):24    for model in models:25        model.eval()26        if use_fp16:27            model.half()28        if use_cuda and not cfg.distributed_training.pipeline_model_parallel:29            model.cuda()30        model.prepare_for_inference_(cfg)31 32 33def construct_transform(patch_image_size):34    mean = [0.5, 0.5, 0.5]35    std = [0.5, 0.5, 0.5]36 37    patch_resize_transform = transforms.Compose([38        lambda image: image.convert("RGB"),39        transforms.Resize((patch_image_size, patch_image_size), interpolation=Image.BICUBIC),40        transforms.ToTensor(),41        transforms.Normalize(mean=mean, std=std),42    ])43 44    return patch_resize_transform45 46 47# Register tasks48tasks.register_task('caption', CaptionTask)49tasks.register_task('refcoco', RefcocoTask)50tasks.register_task('vqa_gen', VqaGenTask)51# turn on cuda if GPU is available52use_cuda = torch.cuda.is_available()53# use fp16 only when GPU is available54use_fp16 = False55 56# download checkpoints57os.system('wget https://ofa-silicon.oss-us-west-1.aliyuncs.com/checkpoints/caption_demo.pt; '58          'mkdir -p checkpoints; mv caption_demo.pt checkpoints/caption_demo.pt')59os.system('wget https://ofa-silicon.oss-us-west-1.aliyuncs.com/checkpoints/refcoco_demo.pt; '60          'mkdir -p checkpoints; mv refcoco_demo.pt checkpoints/refcoco_demo.pt')61os.system('wget https://ofa-silicon.oss-us-west-1.aliyuncs.com/checkpoints/general_demo.pt; '62          'mkdir -p checkpoints; mv general_demo.pt checkpoints/general_demo.pt')63 64# Load ckpt & config for Image Captioning65caption_overrides = {"bpe_dir": "utils/BPE", "eval_cider": False, "beam": 5,66             "max_len_b": 16, "no_repeat_ngram_size": 3, "seed": 7}67caption_models, caption_cfg, caption_task = checkpoint_utils.load_model_ensemble_and_task(68    utils.split_paths('checkpoints/caption_demo.pt'),69    arg_overrides=caption_overrides70)71 72# Load ckpt & config for Refcoco73refcoco_overrides = {"bpe_dir": "utils/BPE", "eval_cider": False, "beam": 5,74             "max_len_b": 16, "no_repeat_ngram_size": 3, "seed": 7}75refcoco_models, refcoco_cfg, refcoco_task = checkpoint_utils.load_model_ensemble_and_task(76    utils.split_paths('checkpoints/refcoco_demo.pt'),77    arg_overrides=refcoco_overrides78)79refcoco_cfg.common.seed = 780refcoco_cfg.generation.beam = 581refcoco_cfg.generation.min_len = 482refcoco_cfg.generation.max_len_a = 083refcoco_cfg.generation.max_len_b = 484refcoco_cfg.generation.no_repeat_ngram_size = 385 86# Load pretrained ckpt & config for VQA87parser = options.get_generation_parser()88input_args = ["", "--task=vqa_gen", "--beam=100", "--unnormalized", "--path=checkpoints/general_demo.pt", "--bpe-dir=utils/BPE"]89args = options.parse_args_and_arch(parser, input_args)90vqa_cfg = convert_namespace_to_omegaconf(args)91vqa_task = tasks.setup_task(vqa_cfg.task)92vqa_models, vqa_cfg = checkpoint_utils.load_model_ensemble(93    utils.split_paths(vqa_cfg.common_eval.path),94    task=vqa_task95)96 97# Load pretrained ckpt & config for Generic Interface98parser = options.get_generation_parser()99input_args = ["", "--task=refcoco", "--beam=10", "--path=checkpoints/general_demo.pt", "--bpe-dir=utils/BPE", "--no-repeat-ngram-size=3", "--patch-image-size=384"]100args = options.parse_args_and_arch(parser, input_args)101general_cfg = convert_namespace_to_omegaconf(args)102general_task = tasks.setup_task(general_cfg.task)103general_models, general_cfg = checkpoint_utils.load_model_ensemble(104    utils.split_paths(general_cfg.common_eval.path),105    task=general_task106)107 108# move models to gpu109move2gpu(caption_models, caption_cfg)110move2gpu(refcoco_models, refcoco_cfg)111move2gpu(vqa_models, vqa_cfg)112move2gpu(general_models, general_cfg)113 114# Initialize generator115caption_generator = caption_task.build_generator(caption_models, caption_cfg.generation)116refcoco_generator = refcoco_task.build_generator(refcoco_models, refcoco_cfg.generation)117vqa_generator = vqa_task.build_generator(vqa_models, vqa_cfg.generation)118vqa_generator.zero_shot = True119vqa_generator.constraint_trie = None120general_generator = general_task.build_generator(general_models, general_cfg.generation)121 122# Construct image transforms123caption_transform = construct_transform(caption_cfg.task.patch_image_size)124refcoco_transform = construct_transform(refcoco_cfg.task.patch_image_size)125vqa_transform = construct_transform(vqa_cfg.task.patch_image_size)126general_transform = construct_transform(general_cfg.task.patch_image_size)127 128# Text preprocess129bos_item = torch.LongTensor([caption_task.src_dict.bos()])130eos_item = torch.LongTensor([caption_task.src_dict.eos()])131pad_idx = caption_task.src_dict.pad()132 133 134def get_symbols_to_strip_from_output(generator):135    if hasattr(generator, "symbols_to_strip_from_output"):136        return generator.symbols_to_strip_from_output137    else:138        return {generator.bos, generator.eos}139 140 141def decode_fn(x, tgt_dict, bpe, generator, tokenizer=None):142    x = tgt_dict.string(x.int().cpu(), extra_symbols_to_ignore=get_symbols_to_strip_from_output(generator))143    token_result = []144    bin_result = []145    img_result = []146    for token in x.strip().split():147      if token.startswith('<bin_'):148        bin_result.append(token)149      elif token.startswith('<code_'):150        img_result.append(token)151      else:152        if bpe is not None:153          token = bpe.decode('{}'.format(token))154        if tokenizer is not None:155          token = tokenizer.decode(token)156        if token.startswith(' ') or len(token_result) == 0:157          token_result.append(token.strip())158        else:159          token_result[-1] += token160 161    return ' '.join(token_result), ' '.join(bin_result), ' '.join(img_result)162 163 164def bin2coord(bins, w_resize_ratio, h_resize_ratio, cfg):165    bin_list = [int(bin[5:-1]) for bin in bins.strip().split()]166    coord_list = []167    coord_list += [bin_list[0] / (cfg.task.num_bins - 1) * cfg.task.max_image_size / w_resize_ratio]168    coord_list += [bin_list[1] / (cfg.task.num_bins - 1) * cfg.task.max_image_size / h_resize_ratio]169    coord_list += [bin_list[2] / (cfg.task.num_bins - 1) * cfg.task.max_image_size / w_resize_ratio]170    coord_list += [bin_list[3] / (cfg.task.num_bins - 1) * cfg.task.max_image_size / h_resize_ratio]171    return coord_list172 173 174def encode_text(text, length=None, append_bos=False, append_eos=False):175    line = [176        caption_task.bpe.encode(' {}'.format(word.strip()))177        if not word.startswith('<code_') and not word.startswith('<bin_') else word178        for word in text.strip().split()179    ]180    line = ' '.join(line)181    s = caption_task.tgt_dict.encode_line(182        line=line,183        add_if_not_exist=False,184        append_eos=False185    ).long()186    if length is not None:187        s = s[:length]188    if append_bos:189        s = torch.cat([bos_item, s])190    if append_eos:191        s = torch.cat([s, eos_item])192    return s193 194 195def construct_sample(image: Image, instruction: str, transform):196    patch_image = transform(image).unsqueeze(0)197    patch_mask = torch.tensor([True])198 199    instruction = encode_text(' {}'.format(instruction.lower().strip()), append_bos=True, append_eos=True).unsqueeze(0)200    instruction_length = torch.LongTensor([s.ne(pad_idx).long().sum() for s in instruction])201    sample = {202        "id": np.array(['42']),203        "net_input": {204            "src_tokens": instruction,205            "src_lengths": instruction_length,206            "patch_images": patch_image,207            "patch_masks": patch_mask,208        }209    }210    return sample211 212 213# Function to turn FP32 to FP16214def apply_half(t):215    if t.dtype is torch.float32:216        return t.to(dtype=torch.half)217    return t218 219 220def inference(image, task_type, instruction):221    if task_type == 'Image Captioning':222        task = caption_task223        models = caption_models224        generator = caption_generator225        instruction = 'what does the image describe?'226        transform = caption_transform227        cfg = caption_cfg228    elif task_type == 'Visual Question Answering':229        task = vqa_task230        models = vqa_models231        generator = vqa_generator232        transform = vqa_transform233        cfg = vqa_cfg234    elif task_type == 'Visual Grounding':235        task = refcoco_task236        models = refcoco_models237        generator = refcoco_generator238        instruction = 'which region does the text " {} " describe?'.format(instruction)239        transform = refcoco_transform240        cfg = refcoco_cfg241    elif task_type == 'General':242        task = general_task243        models = general_models244        generator = general_generator245        transform = general_transform246        cfg = general_cfg247    else:248        raise NotImplementedError249 250    # Construct input sample & preprocess for GPU if cuda available251    sample = construct_sample(image, instruction, transform)252    sample = utils.move_to_cuda(sample) if use_cuda else sample253    sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample254 255    # Generate result256    with torch.no_grad():257        hypos = task.inference_step(generator, models, sample)258        tokens, bins, imgs = decode_fn(hypos[0][0]["tokens"], task.tgt_dict, task.bpe, generator)259 260    if bins.strip() != '':261        w, h = image.size262        w_resize_ratio = task.cfg.patch_image_size / w263        h_resize_ratio = task.cfg.patch_image_size / h264        img = np.asarray(image)265        coord_list = bin2coord(bins, w_resize_ratio, h_resize_ratio, cfg)266        cv2.rectangle(267            img,268            (int(coord_list[0]), int(coord_list[1])),269            (int(coord_list[2]), int(coord_list[3])),270            (0, 255, 0),271            3272        )273        return img, None274    else:275        return None, tokens276 277inputs = [gr.inputs.Image(type='pil'), gr.inputs.Radio(choices=['Image Captioning',"Visual Question Answering", "Visual Grounding", "General"], type="value", default="Image Captioning", label="Task"), gr.inputs.Textbox(lines=1, label="Instruction")]278outputs = [gr.outputs.Image(type='pil'), 'text']279examples = [280    ['examples/pokemons.jpeg', 'Image Captioning', None],281    ['examples/cats.jpeg', 'Visual Question Answering', 'where are the cats?'],282    ['examples/one_piece.jpeg', 'Visual Grounding', 'a man in a straw hat and a red dress'],283    ['examples/three_houses.jpeg', 'General', 'which region does the text " a grey car " describe?'],284    ['examples/three_houses.jpeg', 'General', 'what color is the left car?']285]286 287title = "OFA"288description = "Gradio Demo for OFA: Unifying Architectures, Tasks, and Modalities Through a Simple Sequence-to-Sequence Learning Framework"289article = "<p style='text-align: center'><a href='http://arxiv.org/abs/2202.03052' target='_blank'>Paper</a> | <a href='https://github.com/OFA-Sys/OFA' target='_blank'>Github Repo</a></p>"290 291io = gr.Interface(fn=inference, inputs=inputs, outputs=outputs,292                  title=title, description=description, article=article, examples=examples, cache_examples=False)293io.launch()