{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%pip install-r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "PyTorch version 1.6.0 available.\n"
     ]
    }
   ],
   "source": [
    "from IPython.display import clear_output, Image, display\n",
    "import PIL.Image\n",
    "import io\n",
    "import json\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 LxmertForQuestionAnswering, LxmertTokenizer\n",
    "import wget\n",
    "import pickle\n",
    "import os\n",
    "\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",
    "GQA_URL = \"https://raw.githubusercontent.com/airsplay/lxmert/master/data/gqa/trainval_label2ans.json\"\n",
    "VQA_URL = \"https://raw.githubusercontent.com/airsplay/lxmert/master/data/vqa/trainval_label2ans.json\"\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()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load object, attribute, and answer labels\n",
    "\n",
    "objids = utils.get_data(OBJ_URL)\n",
    "attrids = utils.get_data(ATTR_URL)\n",
    "gqa_answers = utils.get_data(GQA_URL)\n",
    "vqa_answers = utils.get_data(VQA_URL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "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/eltoto/.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"
     ]
    }
   ],
   "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",
    "lxmert_tokenizer = LxmertTokenizer.from_pretrained(\"unc-nlp/lxmert-base-uncased\")\n",
    "lxmert_gqa = LxmertForQuestionAnswering.from_pretrained(\"unc-nlp/lxmert-gqa-uncased\")\n",
    "lxmert_vqa = LxmertForQuestionAnswering.from_pretrained(\"unc-nlp/lxmert-vqa-uncased\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 80140 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: ['Where is the cat?']\n",
      "prediction from LXMERT GQA: desk\n",
      "prediction from LXMERT VQA: desk\n",
      "Question: ['What is near the disk?']\n",
      "prediction from LXMERT GQA: can\n",
      "prediction from LXMERT VQA: cat\n",
      "Question: ['What is the color of the table?']\n",
      "prediction from LXMERT GQA: brown\n",
      "prediction from LXMERT VQA: brown\n",
      "Question: ['What is the color of the cat?']\n",
      "prediction from LXMERT GQA: black\n",
      "prediction from LXMERT VQA: black and white\n",
      "Question: ['What is the shape of the monitor?']\n",
      "prediction from LXMERT GQA: square\n",
      "prediction from LXMERT VQA: rectangle\n"
     ]
    }
   ],
   "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\")\n",
    "\n",
    "for test_question in test_questions_for_url2:\n",
    "    # run lxmert\n",
    "    test_question = [test_question]\n",
    "\n",
    "    inputs = lxmert_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",
    "    # run lxmert(s)\n",
    "    output_gqa = lxmert_gqa(\n",
    "        input_ids=inputs.input_ids,\n",
    "        attention_mask=inputs.attention_mask,\n",
    "        visual_feats=features,\n",
    "        visual_pos=normalized_boxes,\n",
    "        token_type_ids=inputs.token_type_ids,\n",
    "        output_attentions=False,\n",
    "    )\n",
    "    output_vqa = lxmert_vqa(\n",
    "        input_ids=inputs.input_ids,\n",
    "        attention_mask=inputs.attention_mask,\n",
    "        visual_feats=features,\n",
    "        visual_pos=normalized_boxes,\n",
    "        token_type_ids=inputs.token_type_ids,\n",
    "        output_attentions=False,\n",
    "    )\n",
    "    # get prediction\n",
    "    pred_vqa = output_vqa[\"question_answering_score\"].argmax(-1)\n",
    "    pred_gqa = output_gqa[\"question_answering_score\"].argmax(-1)\n",
    "    print(\"Question:\", test_question)\n",
    "    print(\"prediction from LXMERT GQA:\", gqa_answers[pred_gqa])\n",
    "    print(\"prediction from LXMERT VQA:\", vqa_answers[pred_vqa])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}