SubashSK777/Visual-Question-Answering

1import os2 3os.system('cd fairseq;'4          'pip install ./; cd ..')5os.system('ls -l')6 7import torch8import numpy as np9import re10from fairseq import utils,tasks11from fairseq import checkpoint_utils12from fairseq import distributed_utils, options, tasks, utils13from fairseq.dataclass.utils import convert_namespace_to_omegaconf14from utils.zero_shot_utils import zero_shot_step15from tasks.mm_tasks.vqa_gen import VqaGenTask16from models.ofa import OFAModel17from PIL import Image18from torchvision import transforms19import gradio as gr20 21# Register VQA task22tasks.register_task('vqa_gen',VqaGenTask)23# turn on cuda if GPU is available24use_cuda = torch.cuda.is_available()25# use fp16 only when GPU is available26use_fp16 = False27 28os.system('wget https://ofa-silicon.oss-us-west-1.aliyuncs.com/checkpoints/ofa_large_384.pt; '29          'mkdir -p checkpoints; mv ofa_large_384.pt checkpoints/ofa_large_384.pt')30 31# specify some options for evaluation32parser = options.get_generation_parser()33input_args = ["", "--task=vqa_gen", "--beam=100", "--unnormalized", "--path=checkpoints/ofa_large_384.pt", "--bpe-dir=utils/BPE"]34args = options.parse_args_and_arch(parser, input_args)35cfg = convert_namespace_to_omegaconf(args)36 37# Load pretrained ckpt & config38task = tasks.setup_task(cfg.task)39models, cfg = checkpoint_utils.load_model_ensemble(40    utils.split_paths(cfg.common_eval.path),41    task=task42)43 44# Move models to GPU45for model in models:46    model.eval()47    if use_fp16:48        model.half()49    if use_cuda and not cfg.distributed_training.pipeline_model_parallel:50        model.cuda()51    model.prepare_for_inference_(cfg)52 53# Initialize generator54generator = task.build_generator(models, cfg.generation)55 56# Image transform57from torchvision import transforms58mean = [0.5, 0.5, 0.5]59std = [0.5, 0.5, 0.5]60 61patch_resize_transform = transforms.Compose([62    lambda image: image.convert("RGB"),63    transforms.Resize((cfg.task.patch_image_size, cfg.task.patch_image_size), interpolation=Image.BICUBIC),64    transforms.ToTensor(),65    transforms.Normalize(mean=mean, std=std),66])67 68# Text preprocess69bos_item = torch.LongTensor([task.src_dict.bos()])70eos_item = torch.LongTensor([task.src_dict.eos()])71pad_idx = task.src_dict.pad()72 73# Normalize the question74def pre_question(question, max_ques_words):75    question = question.lower().lstrip(",.!?*#:;~").replace('-', ' ').replace('/', ' ')76    question = re.sub(77        r"\s{2,}",78        ' ',79        question,80    )81    question = question.rstrip('\n')82    question = question.strip(' ')83    # truncate question84    question_words = question.split(' ')85    if len(question_words) > max_ques_words:86        question = ' '.join(question_words[:max_ques_words])87    return question88 89def encode_text(text, length=None, append_bos=False, append_eos=False):90    s = task.tgt_dict.encode_line(91        line=task.bpe.encode(text),92        add_if_not_exist=False,93        append_eos=False94    ).long()95    if length is not None:96        s = s[:length]97    if append_bos:98        s = torch.cat([bos_item, s])99    if append_eos:100        s = torch.cat([s, eos_item])101    return s102 103# Construct input for open-domain VQA task104def construct_sample(image: Image, question: str):105    patch_image = patch_resize_transform(image).unsqueeze(0)106    patch_mask = torch.tensor([True])107 108    question = pre_question(question, task.cfg.max_src_length)109    question = question + '?' if not question.endswith('?') else question110    src_text = encode_text(' {}'.format(question), append_bos=True, append_eos=True).unsqueeze(0)111 112    src_length = torch.LongTensor([s.ne(pad_idx).long().sum() for s in src_text])113    ref_dict = np.array([{'yes': 1.0}]) # just placeholder114    sample = {115        "id":np.array(['42']),116        "net_input": {117            "src_tokens": src_text,118            "src_lengths": src_length,119            "patch_images": patch_image,120            "patch_masks": patch_mask,121        },122        "ref_dict": ref_dict,123    }124    return sample125  126# Function to turn FP32 to FP16127def apply_half(t):128    if t.dtype is torch.float32:129        return t.to(dtype=torch.half)130    return t131 132 133# Function for image captioning134def open_domain_vqa(Image, Question):135    sample = construct_sample(Image, Question)136    sample = utils.move_to_cuda(sample) if use_cuda else sample137    sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample138    # Run eval step for open-domain VQA139    with torch.no_grad():140        result, scores = zero_shot_step(task, generator, models, sample)141    return result[0]['answer']142 143 144title = "OFA-Visual_Question_Answering"145description = "Gradio Demo for OFA-Visual_Question_Answering. Upload your own image (high-resolution images are recommended) or click any one of the examples, and click " \146              "\"Submit\" and then wait for OFA's answer. "147article = "<p style='text-align: center'><a href='https://github.com/OFA-Sys/OFA' target='_blank'>OFA Github " \148          "Repo</a></p> "149examples = [['cat-4894153_1920.jpg', 'where are the cats?'], ['men-6245003_1920.jpg', 'how many people are in the image?'], ['labrador-retriever-7004193_1920.jpg', 'what breed is the dog in the picture?'], ['Starry_Night.jpeg', 'what style does the picture belong to?']]150io = gr.Interface(fn=open_domain_vqa, inputs=[gr.inputs.Image(type='pil'), "textbox"], outputs=gr.outputs.Textbox(label="Answer"),151                  title=title, description=description, article=article, examples=examples,152                  allow_flagging=False, allow_screenshot=False)153io.launch()