{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "# %pip install-r requirements.txt"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "**Note**: This demo is adapted from the LXMERT Demo present here: https://github.com/huggingface/transformers/tree/main/examples/research_projects/lxmert"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "from IPython.display import Image, display\n",
    "import PIL.Image\n",
    "import io\n",
    "import torch\n",
    "import numpy as np\n",
    "from processing_image import Preprocess\n",
    "from visualizing_image import SingleImageViz\n",
    "from modeling_frcnn import GeneralizedRCNN\n",
    "from utils import Config\n",
    "import utils\n",
    "from transformers import VisualBertForQuestionAnswering, BertTokenizerFast\n",
    "\n",
    "# URL = \"https://raw.githubusercontent.com/airsplay/py-bottom-up-attention/master/demo/data/images/input.jpg\"\n",
    "URL = \"https://vqa.cloudcv.org/media/test2014/COCO_test2014_000000262567.jpg\"\n",
    "OBJ_URL = \"https://raw.githubusercontent.com/airsplay/py-bottom-up-attention/master/demo/data/genome/1600-400-20/objects_vocab.txt\"\n",
    "ATTR_URL = \"https://raw.githubusercontent.com/airsplay/py-bottom-up-attention/master/demo/data/genome/1600-400-20/attributes_vocab.txt\"\n",
    "VQA_URL = \"https://dl.fbaipublicfiles.com/pythia/data/answers_vqa.txt\"\n",
    "\n",
    "\n",
    "# for visualizing output\n",
    "def showarray(a, fmt=\"jpeg\"):\n",
    "    a = np.uint8(np.clip(a, 0, 255))\n",
    "    f = io.BytesIO()\n",
    "    PIL.Image.fromarray(a).save(f, fmt)\n",
    "    display(Image(data=f.getvalue()))"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "2021-08-11 04:32:30.532299: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "# load object, attribute, and answer labels\n",
    "\n",
    "objids = utils.get_data(OBJ_URL)\n",
    "attrids = utils.get_data(ATTR_URL)\n",
    "vqa_answers = utils.get_data(VQA_URL)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "# load models and model components\n",
    "frcnn_cfg = Config.from_pretrained(\"unc-nlp/frcnn-vg-finetuned\")\n",
    "\n",
    "frcnn = GeneralizedRCNN.from_pretrained(\"unc-nlp/frcnn-vg-finetuned\", config=frcnn_cfg)\n",
    "\n",
    "image_preprocess = Preprocess(frcnn_cfg)\n",
    "\n",
    "bert_tokenizer = BertTokenizerFast.from_pretrained(\"bert-base-uncased\")\n",
    "visualbert_vqa = VisualBertForQuestionAnswering.from_pretrained(\"uclanlp/visualbert-vqa\")"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "loading configuration file cache\n",
      "loading weights file https://cdn.huggingface.co/unc-nlp/frcnn-vg-finetuned/pytorch_model.bin from cache at /home/crocoder/.cache/torch/transformers/57f6df6abe353be2773f2700159c65615babf39ab5b48114d2b49267672ae10f.77b59256a4cf8343ae0f923246a81489fc8d82f98d082edc2d2037c977c0d9d0\n",
      "All model checkpoint weights were used when initializing GeneralizedRCNN.\n",
      "\n",
      "All the weights of GeneralizedRCNN were initialized from the model checkpoint at unc-nlp/frcnn-vg-finetuned.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use GeneralizedRCNN for predictions without further training.\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "# image viz\n",
    "frcnn_visualizer = SingleImageViz(URL, id2obj=objids, id2attr=attrids)\n",
    "# run frcnn\n",
    "images, sizes, scales_yx = image_preprocess(URL)\n",
    "output_dict = frcnn(\n",
    "    images,\n",
    "    sizes,\n",
    "    scales_yx=scales_yx,\n",
    "    padding=\"max_detections\",\n",
    "    max_detections=frcnn_cfg.max_detections,\n",
    "    return_tensors=\"pt\",\n",
    ")\n",
    "# add boxes and labels to the image\n",
    "\n",
    "frcnn_visualizer.draw_boxes(\n",
    "    output_dict.get(\"boxes\"),\n",
    "    output_dict.pop(\"obj_ids\"),\n",
    "    output_dict.pop(\"obj_probs\"),\n",
    "    output_dict.pop(\"attr_ids\"),\n",
    "    output_dict.pop(\"attr_probs\"),\n",
    ")\n",
    "showarray(frcnn_visualizer._get_buffer())"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "/home/crocoder/anaconda3/envs/transformers_env/lib/python3.8/site-packages/torch/nn/functional.py:718: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at  /pytorch/c10/core/TensorImpl.h:1156.)\n",
      "  return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "<IPython.core.display.Image object>"
      ],
      "image/jpeg": "[encoded data, 82072 characters]"
     },
     "metadata": {}
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "# test_questions_for_url1 = [\n",
    "#     \"Where is this scene?\",\n",
    "#     \"what is the man riding?\",\n",
    "#     \"What is the man wearing?\",\n",
    "#     \"What is the color of the horse?\"\n",
    "# ]\n",
    "test_questions_for_url2 = [\n",
    "    \"Where is the cat?\",\n",
    "    \"What is near the disk?\",\n",
    "    \"What is the color of the table?\",\n",
    "    \"What is the color of the cat?\",\n",
    "    \"What is the shape of the monitor?\",\n",
    "]\n",
    "\n",
    "# Very important that the boxes are normalized\n",
    "# normalized_boxes = output_dict.get(\"normalized_boxes\")\n",
    "features = output_dict.get(\"roi_features\")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "source": [
    "for test_question in test_questions_for_url2:\n",
    "    test_question = [test_question]\n",
    "\n",
    "    inputs = bert_tokenizer(\n",
    "        test_question,\n",
    "        padding=\"max_length\",\n",
    "        max_length=20,\n",
    "        truncation=True,\n",
    "        return_token_type_ids=True,\n",
    "        return_attention_mask=True,\n",
    "        add_special_tokens=True,\n",
    "        return_tensors=\"pt\",\n",
    "    )\n",
    "\n",
    "    output_vqa = visualbert_vqa(\n",
    "        input_ids=inputs.input_ids,\n",
    "        attention_mask=inputs.attention_mask,\n",
    "        visual_embeds=features,\n",
    "        visual_attention_mask=torch.ones(features.shape[:-1]),\n",
    "        token_type_ids=inputs.token_type_ids,\n",
    "        output_attentions=False,\n",
    "    )\n",
    "    # get prediction\n",
    "    pred_vqa = output_vqa[\"logits\"].argmax(-1)\n",
    "    print(\"Question:\", test_question)\n",
    "    print(\"prediction from VisualBert VQA:\", vqa_answers[pred_vqa])"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Question: ['Where is the cat?']\n",
      "prediction from VisualBert VQA: outside\n",
      "Question: ['What is near the disk?']\n",
      "prediction from VisualBert VQA: nothing\n",
      "Question: ['What is the color of the table?']\n",
      "prediction from VisualBert VQA: brown\n",
      "Question: ['What is the color of the cat?']\n",
      "prediction from VisualBert VQA: gray\n",
      "Question: ['What is the shape of the monitor?']\n",
      "prediction from VisualBert VQA: square\n"
     ]
    }
   ],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.0 64-bit ('transformers_env': conda)"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  },
  "interpreter": {
   "hash": "f237d186bbb22b392353378fb98a8d08e33f23f14150c8880e3780871939e71d"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}