Spaces:
Runtime error
Runtime error
| # Requirements.txt | |
| from torch import cuda | |
| from transformers import T5Tokenizer, T5ForConditionalGeneration | |
| import gradio as gr | |
| from utils import convert_ans_to_token, convert_ques_to_token, rotate, convert_token_to_ques, convert_token_to_answer | |
| from modeling import LaTr_for_pretraining, LaTr_for_finetuning, LaTrForVQA | |
| from dataset import load_json_file, get_specific_file, resize_align_bbox, get_tokens_with_boxes, create_features | |
| import torch.nn as nn | |
| from PIL import Image, ImageDraw | |
| import pytesseract | |
| from tqdm.auto import tqdm | |
| import numpy as np | |
| import json | |
| import os | |
| import torch | |
| from torchvision import transforms | |
| # install PyTesseract | |
| os.system('pip install -q pytesseract') | |
| os.environ["TOKENIZERS_PARALLELISM"] = "false" | |
| # Default Library import | |
| # Visualization libraries | |
| # Specific libraries of LaTr | |
| # Setting the hyperparameters as well as primary configurations | |
| PAD_TOKEN_BOX = [0, 0, 0, 0] | |
| max_seq_len = 512 | |
| batch_size = 2 | |
| target_size = (500, 384) | |
| t5_model = "t5-base" | |
| device = 'cuda' if cuda.is_available() else 'cpu' | |
| # Configuration for the model | |
| config = { | |
| 't5_model': 't5-base', | |
| 'vocab_size': 32128, | |
| 'hidden_state': 768, | |
| 'max_2d_position_embeddings': 1001, | |
| 'classes': 32128, # number of tokens | |
| 'seq_len': 512 | |
| } | |
| tokenizer = T5Tokenizer.from_pretrained(t5_model) | |
| latr = LaTrForVQA(config) | |
| url = 'https://www.kaggleusercontent.com/kf/99663112/eyJhbGciOiJkaXIiLCJlbmMiOiJBMTI4Q0JDLUhTMjU2In0..2HGa6jqeAbugMJYxSkh7eA.XkaLSf8XlITet17Bscupegw9zWLw-IEizSy1lM-_PJF_Gfj-YuinOpDw4ad0M8r-s3WlnclQhHYrd2seaZVjBmkm5WSE6Dae1fW54dnNhyWF5w5O2VafNar7QSuUTSRzacJcmtqI1ypL3OZofwXuETbXq4weeqfDptFS5luxuV0P4Vaer_xEgfsdld6v8O5jjMXwb1CVmPCjMdZUE-HTgzTDiwv3Lb-P3dkRgU7q-iI5GeYZCODYGrX-koxya9DlfzKQZXmJmvtMj45vUZ8OSRB0_hTc7UosQanA-SalWznnOuyOgwl4hMag5toTomriWsxfvJIRBn9CYgFcvUJNqO_kDzBUoAwnagjcxXeEIJTJglwAl9Rs37XyfJAZr7yQ_YTXeRW1j2QMsT_M3qtS96IKRTpsqPVibl8Vrs9Q5g_vKccIQR9t7R9ma_DZLwjWYhDvDO06AZqtdaYGfWaOrbqe8dDvJkZoHsZEO8ukpIH6YNLyCO_dqgRsE77I9jqxiUqQh1KnuNv2hGRSlQR7u8OF7lpiRS7JEwj2MaxlzD58dyhOOLDqrbLp7XWrgV79EQcRYHFSMfhDvG0zmGvHjWGAg-LGhnYIc0NMVhyRv5Pfta9WYEl4qXxCTZWe4olgV79WHLqksQMVyTteheB36n4biHZKx4KZj7k-j3aSI72DIAvj7_UFeHxUTTZ1c6MB.7BF6J5MPMuhQFU48xVZ2qQ/models/epoch=0-step=34602.ckpt' | |
| try: | |
| latr = latr.load_from_checkpoint(url) | |
| print("Checkpoint loaded successfully") | |
| except: | |
| print("Checkpoint not loaded") | |
| pass | |
| image = gr.inputs.Image(type="pil") | |
| question = gr.inputs.Textbox(label="Question") | |
| answer = gr.outputs.Textbox(label="Predicted answer") | |
| examples = [["remote.jpg", "what number is the button near the top left?"]] | |
| from transformers import ViTFeatureExtractor, ViTModel | |
| vit_feat_extract = ViTFeatureExtractor("google/vit-base-patch16-224-in21k") | |
| import torchvision | |
| import numpy as np | |
| def answer_question(image, question): | |
| # Extracting features from the image | |
| image.save("sample.png") | |
| img, boxes, tokenized_words = create_features("sample.png", | |
| tokenizer=tokenizer, | |
| target_size=target_size, | |
| max_seq_length=max_seq_len, | |
| use_ocr=True | |
| ) | |
| ## Converting the boxes as per the format required for model input | |
| boxes = torch.as_tensor(boxes, dtype=torch.int32) | |
| width = (boxes[:, 2] - boxes[:, 0]).view(-1, 1) | |
| height = (boxes[:, 3] - boxes[:, 1]).view(-1, 1) | |
| boxes = torch.cat([boxes, width, height], axis = -1) | |
| ## Clamping the value,as some of the box values are out of bound | |
| boxes[:, 0] = torch.clamp(boxes[:, 0], min = 0, max = 0) | |
| boxes[:, 2] = torch.clamp(boxes[:, 2], min = 1000, max = 1000) | |
| boxes[:, 4] = torch.clamp(boxes[:, 4], min = 1000, max = 1000) | |
| boxes[:, 1] = torch.clamp(boxes[:, 1], min = 0, max = 0) | |
| boxes[:, 3] = torch.clamp(boxes[:, 3], min = 1000, max = 1000) | |
| boxes[:, 5] = torch.clamp(boxes[:, 5], min = 1000, max = 1000) | |
| ## Tensor tokenized words | |
| tokenized_words = torch.as_tensor(tokenized_words, dtype=torch.int32) | |
| img = np.array(img) | |
| img = torchvision.transforms.ToTensor()(img) | |
| question = convert_ques_to_token(question = question, tokenizer = tokenizer) | |
| ## Expanding the dimension for inference | |
| boxes = boxes.unsqueeze(0) | |
| tokenized_words = tokenized_words.unsqueeze(0) | |
| question = question.unsqueeze(0) | |
| # print("Shape of Image is:", img.shape) | |
| img = vit_feat_extract(img, return_tensors = 'pt')['pixel_values'] | |
| if int(len(img.shape)) == 3: | |
| img = img.unsqueeze(0) | |
| encoding = {'img': img, 'boxes': boxes, 'tokenized_words': tokenized_words, 'question': question} | |
| with torch.no_grad(): | |
| logits = latr.forward(encoding) | |
| logits = logits.squeeze(0) | |
| _, preds = torch.max(logits, dim = 1) | |
| preds = preds.detach().cpu() | |
| mask = torch.clamp(preds, min = 0, max = 1) | |
| last_non_zero_argument = (mask != 0).nonzero()[1][-1] | |
| predicted_ans = convert_token_to_ques(preds[:last_non_zero_argument], tokenizer) | |
| return predicted_ans | |
| # Taken from here: https://huggingface.co/spaces/nielsr/vilt-vqa/blob/main/app.py | |
| title = "Interactive demo: LaTr (Layout Aware Transformer) for VQA" | |
| description = "Gradio Demo for LaTr (Layout Aware Transformer),trained on TextVQA Dataset. To use it, simply upload your image and type a question and click 'submit', or click one of the examples to load them. Read more at the links below." | |
| article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2112.12494' target='_blank'>LaTr: Layout-aware transformer for scene-text VQA,a novel multimodal architecture for Scene Text Visual Question Answering (STVQA)</a> | <a href='https://github.com/uakarsh/latr' target='_blank'>Github Repo</a></p>" | |
| examples = [['remote.png', "Is remote present in the picture?"]] | |
| interface = gr.Interface(fn=answer_question, | |
| inputs=[image, question], | |
| outputs=answer, | |
| examples=examples, | |
| title=title, | |
| description=description, | |
| article=article, | |
| enable_queue=True) | |
| interface.launch(debug=True) | |