multimodal/examples/mdetr/MDETRTutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9ae1d1ec-663d-4ce3-bb37-d5abb7a00e0f",
   "metadata": {},
   "source": [
    "# Multimodal detection with MDETR in TorchMultimodal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36aa5885-2cdb-4adb-889c-d605e2f3ee93",
   "metadata": {},
   "source": [
    "[MDETR](https://arxiv.org/pdf/2104.12763.pdf) (Kamath et al.) is a multimodal detector that detects objects in an image conditioned on text queries. It is an extension of the [DETR](https://arxiv.org/pdf/2005.12872.pdf) model, which treats object detection as a set prediction problem, performing a bipartite matching between the ground truth and predicted bounding boxes. MDETR extends this approach to vision & language tasks with the addition of a contrastive alignment loss and a soft token prediction loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c72b4ce-197d-4bdb-a0f4-0e8ceae8f7f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image\n",
    "from IPython.core.display import HTML \n",
    "Image(url=\"https://production-media.paperswithcode.com/methods/Screen_Shot_2021-08-11_at_10.03.50_AM.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "204615b1-dc55-4037-b571-5e1247fd4e18",
   "metadata": {},
   "source": [
    "MDETR shows strong performance on a number of downstream tasks, including visual question answering (VQA), phrase grounding, and referring expression comprehension and segmentation. In this tutorial, we first give an introduction to the components of MDETR in TorchMultimodal, then provide a walkthrough of how to apply MDETR on both phrase grounding and VQA tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d2a2708-836f-42f0-a43c-2b768e098300",
   "metadata": {},
   "source": [
    "## MDETR in TorchMultimodal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49bdd5d6-5a6e-4a11-a532-3a2af24168f6",
   "metadata": {},
   "source": [
    "TorchMultimodal provides access to MDETR components at multiple different levels of granularity. This allows users to easily plug in their own modules for experimentation without needing to change the underlying implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd93cf74-7d5a-411f-af49-ea81c0085d79",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simplest entry point: model builder function.\n",
    "# Provides prebuilt model from the paper with ResNet101 backbone.\n",
    "from torchmultimodal.models.mdetr.model import mdetr_resnet101\n",
    "mdetr_rn101 = mdetr_resnet101()\n",
    "\n",
    "# The builder function contains sensible default values but can also be easily modified.\n",
    "# E.g. we can modify the number of object queries and change the transformer FFN dimension\n",
    "modified_mdetr_rn101 = mdetr_resnet101(num_queries=100, transformer_dim_feedforward=1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e131192-23b1-4a73-af32-05a8251470f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# mdetr_resnet101 is an instantiation of the MDETR model class.\n",
    "# This class is composed of modular components which can be swapped out for easy experimentation.\n",
    "from torchmultimodal.models.mdetr.model import MDETR\n",
    "\n",
    "# Other imports we will use\n",
    "import math\n",
    "from torch import nn\n",
    "from torchmultimodal.modules.layers.mlp import MLP\n",
    "from torchmultimodal.models.mdetr.model import FeatureResizer\n",
    "from torchmultimodal.models.mdetr.image_encoder import mdetr_resnet101_backbone, MaskedIntermediateLayer, PositionEmbedding2D\n",
    "from torchmultimodal.models.mdetr.text_encoder import mdetr_roberta_text_encoder\n",
    "from torchmultimodal.models.mdetr.transformer import mdetr_transformer\n",
    "\n",
    "# Set some dimensions as constant\n",
    "image_embedding_dim = 2048\n",
    "text_embedding_dim = 768\n",
    "transformer_dim = 256\n",
    "num_queries = 100\n",
    "num_classes = 255\n",
    "\n",
    "# The core components of the MDETR model class are:\n",
    "#   - image backbone \n",
    "#   - text encoder \n",
    "#   - multimodal transformer\n",
    "\n",
    "# Encoder and transformer components also have builder functions.\n",
    "image_backbone = mdetr_resnet101_backbone()\n",
    "text_encoder = mdetr_roberta_text_encoder()\n",
    "transformer = mdetr_transformer()\n",
    "\n",
    "\n",
    "# Other components include\n",
    "#   - positional embeddings: 2D embeddings added to image backbone outputs\n",
    "#   - text/image projections (between unimodal encoders and multimodal transformer)\n",
    "#   - query embeddings: learned object query embeddings, used in multimodal transformer decoder\n",
    "#   - bounding box embeddings: map multimodal transformer outputs to bounding box coordinates\n",
    "#   - class embeddings: map multimodal transformer outputs to class logits \n",
    "\n",
    "pos_embed = PositionEmbedding2D(num_pos_feats=128, scale=2 * math.pi)\n",
    "text_projection = FeatureResizer(text_embedding_dim, transformer_dim)\n",
    "image_projection = nn.Conv2d(image_embedding_dim, transformer_dim, kernel_size=1)\n",
    "query_embed = nn.Embedding(num_queries, transformer_dim)\n",
    "bbox_embed = MLP(transformer_dim, 4, [transformer_dim] * 2, dropout=0.0)\n",
    "class_embed = nn.Linear(transformer_dim, num_classes + 1)\n",
    "\n",
    "# This will give us the original ResNet101 MDETR model\n",
    "mdetr = MDETR(\n",
    "    image_backbone = image_backbone,\n",
    "    text_encoder = text_encoder, \n",
    "    transformer = transformer,\n",
    "    pos_embed = pos_embed, \n",
    "    text_projection = text_projection,\n",
    "    image_projection = image_projection, \n",
    "    query_embed = query_embed,\n",
    "    bbox_embed = bbox_embed,\n",
    "    class_embed = class_embed,\n",
    ")\n",
    "\n",
    "# Swap out text encoder and class embedding for other components\n",
    "from transformers import DistilBertModel\n",
    "distilbert_text_encoder = DistilBertModel.from_pretrained(\"distilbert-base-uncased\")\n",
    "mlp_class_embed = MLP(transformer_dim, num_classes + 1, [transformer_dim] * 2, activation=nn.GELU, dropout=0.1)\n",
    "modified_mdetr = MDETR(\n",
    "    image_backbone = image_backbone,\n",
    "    text_encoder = distilbert_text_encoder, \n",
    "    transformer = transformer,\n",
    "    pos_embed = pos_embed, \n",
    "    text_projection = text_projection,\n",
    "    image_projection = image_projection, \n",
    "    query_embed = query_embed,\n",
    "    bbox_embed = bbox_embed,\n",
    "    class_embed = class_embed,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95e99f7c-ca24-4ea4-b261-d4331f3a7733",
   "metadata": {},
   "source": [
    "## Phrase grounding"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd17c053-09db-43d3-8799-1c35082521b0",
   "metadata": {},
   "source": [
    "This section demonstrates how to use MDETR for [phrase grounding](https://paperswithcode.com/task/phrase-grounding/codeless). In phrase grounding, the goal is to associate noun phrases in the caption of an `(image, text)` pair to regions in the image. This is in line with MDETR's soft token alignment loss, which means we can perform inference on this task directly from the pretrained checkpoint. Here we will use a few samples from the [Flickr30k](https://paperswithcode.com/dataset/flickr30k#:~:text=Introduced%20by%20Young%20et%20al,sentences%20provided%20by%20human%20annotators.) validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0456d77d-e563-4e54-8995-754512daacee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TorchMultimodal also provides entry points to apply MDETR on specific tasks.\n",
    "# These wrap the MDETR class from the preceding section with task-specific losses and/or heads.\n",
    "import torch\n",
    "from torchmultimodal.models.mdetr.model import mdetr_for_phrase_grounding\n",
    "\n",
    "mdetr_pg = mdetr_for_phrase_grounding()\n",
    "checkpoint = torch.hub.load_state_dict_from_url(\"https://pytorch.s3.amazonaws.com/models/multimodal/mdetr/pretrained_resnet101_checkpoint.pth\")\n",
    "mdetr_pg.load_state_dict(checkpoint[\"model_ema\"])\n",
    "mdetr_pg.eval()\n",
    "print(\"done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5806d903-96b3-4f39-86cf-4953d4fbbd6c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Download a few sample images and corresponding annotations\n",
    "import os\n",
    "asset_dir = \"assets/phrase_grounding\"\n",
    "img_dir = os.path.join(asset_dir, \"images\")\n",
    "os.makedirs(img_dir, exist_ok=True)\n",
    "\n",
    "!wget -nc \"https://pytorch.s3.amazonaws.com/torchmultimodal/examples/mdetr/assets/phrase_grounding/images/4287766241.jpg\" -P $img_dir\n",
    "!wget -nc \"https://pytorch.s3.amazonaws.com/torchmultimodal/examples/mdetr/assets/phrase_grounding/images/4591476340.jpg\" -P $img_dir\n",
    "!wget -nc \"https://pytorch.s3.amazonaws.com/torchmultimodal/examples/mdetr/assets/phrase_grounding/flickr30k_sample_annotations.json\" -P $asset_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e3ade3a-b3bd-4ec1-86fb-46b1bd24676a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from examples.mdetr.data.dataset import ModulatedDetection\n",
    "import torchvision.transforms as T\n",
    "from transformers import RobertaTokenizerFast\n",
    "\n",
    "# Construct tokenizer, image transform, and dataset on sample data \n",
    "# We can also pass transforms directly to the dataset class via MDETRTransform. \n",
    "# We opt not to so that we can inspect each sample first.\n",
    "tokenizer = RobertaTokenizerFast.from_pretrained(\"roberta-base\")\n",
    "img_transform = T.Compose([\n",
    "    T.Resize(800),\n",
    "    T.ToTensor(),\n",
    "    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "])\n",
    "dataset = ModulatedDetection(\n",
    "        img_dir,\n",
    "        os.path.join(asset_dir, \"flickr30k_sample_annotations.json\"),\n",
    "        transforms=None,\n",
    "        return_tokens=True,\n",
    "        tokenizer=tokenizer,\n",
    "        is_train=False,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcd4df2e-912a-4e8e-8bcc-f62b892f185b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define some useful utilities\n",
    "from collections import defaultdict\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from torchvision.ops.boxes import box_convert\n",
    "\n",
    "COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],\n",
    "          [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]\n",
    "\n",
    "# Util to rescale predicted boxes to match image size\n",
    "def rescale_boxes(boxes, size):\n",
    "    img_w, img_h = size\n",
    "    b = box_convert(boxes, \"cxcywh\", \"xyxy\")\n",
    "    b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)\n",
    "    return b\n",
    "\n",
    "# Plotting util function\n",
    "def make_plot(pil_img, scores, boxes, labels):\n",
    "    plt.figure(figsize=(16,10))\n",
    "    np_image = np.array(pil_img)\n",
    "    ax = plt.gca()\n",
    "    colors = COLORS * 100\n",
    "    assert len(scores) == len(boxes) == len(labels)\n",
    "    for s, box, l, c in zip(scores, boxes, labels, colors):\n",
    "        (xmin, ymin, xmax, ymax) = box.detach().numpy()\n",
    "        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,\n",
    "                                   fill=False, color=c, linewidth=3))\n",
    "        text = f'{l}: {s:0.2f}'\n",
    "        ax.text(xmin, ymin, text, fontsize=15, bbox=dict(facecolor='white', alpha=0.8))\n",
    "\n",
    "    plt.imshow(np_image)\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "\n",
    "# Postprocessing to rescale boxes and extract predicted spans of tokens\n",
    "def postprocess_for_plot(outputs, text, tokenized, img_size, do_vqa=False):        \n",
    "    model_outputs = outputs.model_output\n",
    "    probs = 1 - model_outputs.pred_logits.softmax(-1)[0, :, -1]\n",
    "    keep = (probs > 0.7)\n",
    "\n",
    "    # convert boxes from [0; 1] to image scales\n",
    "    boxes_scaled = rescale_boxes(model_outputs.pred_boxes[0, keep], img_size)\n",
    "\n",
    "    # Extract the text spans predicted by each box\n",
    "    positive_tokens = (model_outputs.pred_logits[0, keep].softmax(-1) > 0.1).nonzero().tolist()\n",
    "    predicted_spans = defaultdict(str)\n",
    "    for tok in positive_tokens:\n",
    "        item, pos = tok\n",
    "        if pos < 255:\n",
    "            span = tokenized.token_to_chars(0, pos)\n",
    "            predicted_spans[item] += \" \" + text[span.start:span.end]\n",
    "    boxes_scaled = [boxes_scaled[int(k)] for k in sorted(list(predicted_spans.keys()))]\n",
    "    labels = [predicted_spans[k] for k in sorted(list(predicted_spans.keys()))]\n",
    "    \n",
    "    return probs[keep], boxes_scaled, labels\n",
    "\n",
    "# Main method to plot results of MDETR inference\n",
    "def plot_inference_results(img, text, model=mdetr_pg, img_transform=img_transform, tokenizer=tokenizer, do_vqa=False):\n",
    "    img_size = img.size\n",
    "    # caption = text[\"caption\"]\n",
    "    tokenized = tokenizer.batch_encode_plus([text], padding=\"longest\", return_tensors=\"pt\")\n",
    "    img_transformed = img_transform(img)\n",
    "    outputs = model([img_transformed], tokenized[\"input_ids\"])\n",
    "    probs, boxes, labels = postprocess_for_plot(outputs, text, tokenized, img_size)\n",
    "    make_plot(img, probs, boxes, labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "472d253b-0db1-4695-be8a-9365c9cb1790",
   "metadata": {},
   "outputs": [],
   "source": [
    "img, target = dataset[0]\n",
    "text = target[\"caption\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fd0f230-6746-4d39-9105-1dea926b908b",
   "metadata": {},
   "outputs": [],
   "source": [
    "img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f459821-b1e2-4da6-b260-d47da4cefd00",
   "metadata": {},
   "outputs": [],
   "source": [
    "text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a430f8b6-4041-4a23-99cb-c8326f4a776c",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_inference_results(img, text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f573692c-6de8-4d3b-8ae9-ec9cb29a7a29",
   "metadata": {},
   "source": [
    "Note that the single phrase \"a couple\" yields several bounding boxes. More discussion on this point can be found in Appendix D of the [MDETR paper](https://arxiv.org/pdf/2104.12763.pdf). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edf12365-5374-4952-b15d-dd41013f060e",
   "metadata": {},
   "source": [
    "We can also pass custom prompts to see how well the model can detect other entities in the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c819d768-2f4f-4391-aa2c-b9c1885319a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_inference_results(img, \"Two lanterns\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "333a38cf-2f29-4f01-b038-82b4a52d077d",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_inference_results(img, \"A backpack and a blue railing\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1f9c868-9c0a-43a2-a254-6c2219abd9e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_inference_results(img, \"A singer with a microphone\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88fa4a19-7157-44e7-9b02-3c407974581f",
   "metadata": {},
   "outputs": [],
   "source": [
    "img, target = dataset[1]\n",
    "text = target[\"caption\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "958e0e07-0137-428d-985c-b78201c3a816",
   "metadata": {},
   "outputs": [],
   "source": [
    "img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19a99a49-fb9d-4fc1-a672-40ca1bd57c11",
   "metadata": {},
   "outputs": [],
   "source": [
    "text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e3d5831-273a-44a4-937d-37aa3217c4ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_inference_results(img, text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ea76b83-260e-4375-8471-0f3800d5f937",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_inference_results(img, \"a white hat, trash bins, a pink headband, and blinds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efff6796-3649-44b8-9cc3-27fde3f660ad",
   "metadata": {},
   "source": [
    "Note that the model will often output predictions even for objects not found in the image. As explained in the [demo](https://colab.research.google.com/github/ashkamath/mdetr/blob/colab/notebooks/MDETR_demo.ipynb) based on the original paper's repo, this is due to the lack of negatives in the training set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "647c1167-21fb-4d77-b205-de5df6759654",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_inference_results(img, \"A singer with a microphone\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d576799a-da0a-4de1-abf1-f68adf951dd3",
   "metadata": {},
   "source": [
    "## Visual Question Answering"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0651238-5418-44f0-8039-1dae483133cd",
   "metadata": {},
   "source": [
    "During pretraining, MDETR uses a fixed number of object queries to detect different objects in the image. (The number of object queries provides an upper bound on the number of objects MDETR can detect.) For visual question answering, the object queries are supplemented by additional queries corresponding to each type of question (as well as one final query used for the question type itself). The resulting embeddings are fed to classification heads which perform predictions.\n",
    "\n",
    "Here we demonstrate inference for a couple of samples from the [GQA dataset](https://cs.stanford.edu/people/dorarad/gqa/index.html). For GQA, the question types are based on GQA's set of [semantic types](https://cs.stanford.edu/people/dorarad/gqa/vis.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ff0ab16-32df-47cb-b6e3-ceb9d4956473",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TorchMultimodal also provides a builder for VQA tasks.\n",
    "from torchmultimodal.models.mdetr.model import mdetr_for_vqa\n",
    "\n",
    "# The default values correspond to the GQA dataset.\n",
    "mdetr_vqa = mdetr_for_vqa()\n",
    "\n",
    "# But to perform VQA on another dataset, one can simply pass a different set of heads, e.g.\n",
    "other_vqa_heads = nn.ModuleDict({\"head1\": nn.Linear(256, 3), \"head2\": nn.Linear(256,12)})\n",
    "mdetr_other_vqa_dataset = mdetr_for_vqa(vqa_heads=other_vqa_heads)\n",
    "\n",
    "\n",
    "checkpoint = torch.hub.load_state_dict_from_url(\n",
    "            \"https://pytorch.s3.amazonaws.com/models/multimodal/mdetr/gqa_resnet101_checkpoint.pth\", map_location=\"cpu\", check_hash=True\n",
    "        )\n",
    "mdetr_vqa.load_state_dict(checkpoint[\"model_ema\"], strict=False)\n",
    "mdetr_vqa.eval()\n",
    "print(\"done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3fd2711-24ac-4399-89e9-5972565a7ad0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download a few sample images and corresponding annotations\n",
    "import os\n",
    "asset_dir = \"assets/vqa\"\n",
    "img_dir = os.path.join(asset_dir, \"images\")\n",
    "os.makedirs(img_dir, exist_ok=True)\n",
    "\n",
    "!wget -nc \"https://pytorch.s3.amazonaws.com/torchmultimodal/examples/mdetr/assets/gqa/images/n216553.jpg\" -P $img_dir\n",
    "!wget -nc \"https://pytorch.s3.amazonaws.com/torchmultimodal/examples/mdetr/assets/gqa/images/n411121.jpg\" -P $img_dir\n",
    "!wget -nc \"https://pytorch.s3.amazonaws.com/torchmultimodal/examples/mdetr/assets/gqa/gqa_sample_annotations.json\" -P $asset_dir\n",
    "!wget -nc \"https://pytorch.s3.amazonaws.com/torchmultimodal/examples/mdetr/assets/gqa/gqa_answer2id.json\" -P $asset_dir\n",
    "!wget -nc \"https://pytorch.s3.amazonaws.com/torchmultimodal/examples/mdetr/assets/gqa/gqa_answer2id_by_type.json\" -P $asset_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e60da55-0af9-4e40-903f-982c89cb2ad5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct GQA dataset. Similar to for phrase grounding, we will perform\n",
    "# the transforms ourselves for more flexible inference. \n",
    "from examples.mdetr.data.dataset import GQADataset\n",
    "from pathlib import Path\n",
    "\n",
    "dataset = GQADataset(\n",
    "            img_dir,\n",
    "            os.path.join(asset_dir, \"gqa_sample_annotations.json\"),\n",
    "            transforms=None,\n",
    "            return_tokens=True,\n",
    "            tokenizer=tokenizer,\n",
    "            ann_folder=Path(asset_dir),\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d389932a-f857-461d-bf1f-1852e0c16e9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define some more utils for plotting VQA predictions\n",
    "\n",
    "# Used for displaying text answers\n",
    "id2answerbytype = {}\n",
    "for ans_type in dataset.answer2id_by_type.keys():                        \n",
    "    curr_reversed_dict = {v: k for k, v in dataset.answer2id_by_type[ans_type].items()}\n",
    "    id2answerbytype[ans_type] = curr_reversed_dict \n",
    "id2answer = {v: k for k, v in dataset.answer2id.items()}\n",
    "types = {v: k for k, v in dataset.type_to_id.items()}\n",
    "    \n",
    "# Get predicted question type and answer along with ground truth\n",
    "def get_question_answer(outputs, answer, types=types, id2answer=id2answer,id2answerbytype=id2answerbytype, top_k=5):\n",
    "    preds = outputs.vqa_preds['answer_type'].softmax(-1)\n",
    "    type_conf, type_pred = preds.max(-1)\n",
    "    for i in range(preds.size(-1)):\n",
    "        print(f\"Question type {types[i]} probability: {round(100.0 * preds[:,i].item(), 2)}\")\n",
    "    type_pred = type_pred.item()\n",
    "    answers = outputs.vqa_preds[f\"answer_{types[type_pred]}\"][0].softmax(-1)\n",
    "    ans_conf, ans = torch.topk(answers, top_k)\n",
    "    predicted_answers = [id2answerbytype[f\"answer_{types[type_pred]}\"][x.item()] for x in ans]\n",
    "    ground_truth = id2answer[answer.item()]\n",
    "    # ground_truth = id2answerbytype[f\"answer_{types[answer_type.item()]}\"][answer.item()]\n",
    "    confidences =  [100 * type_conf.item() * x.item() for x in ans_conf]\n",
    "    return predicted_answers, confidences, ground_truth\n",
    "    \n",
    "# Slight modification to original inference plotting function to also print QA predictions\n",
    "def plot_qa_inference_results(img, question, answer, model=mdetr_vqa, img_transform=img_transform, tokenizer=tokenizer, top_k=5):\n",
    "    img_size = img.size\n",
    "    # caption = text[\"caption\"]\n",
    "    tokenized = tokenizer.batch_encode_plus([question], padding=\"longest\", return_tensors=\"pt\")\n",
    "    img_transformed = img_transform(img)\n",
    "    outputs = model([img_transformed], tokenized[\"input_ids\"])\n",
    "    probs, boxes, labels = postprocess_for_plot(outputs, question, tokenized, img_size)\n",
    "    print(question)\n",
    "    make_plot(img, probs, boxes, labels)\n",
    "    predicted_answers, confidences, ground_truth = get_question_answer(outputs, answer, top_k=top_k)\n",
    "    for i in range(len(predicted_answers)):\n",
    "        print(f\"Predicted answer #{i+1}: {predicted_answers[i]}, confidence: {round(confidences[i], 2)}\")\n",
    "    # print(f\")\n",
    "    # print(f\"Predicted answer: {predicted_answers}\\t confidence={[round(confidences, 2)}%\")\n",
    "    print(f\"Ground truth answer: {ground_truth}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd1716a2-c0a1-4b1c-baa1-069ce51cfe76",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0\n",
    "img = dataset[i][0]\n",
    "question = dataset[i][1][\"caption\"]\n",
    "answer = dataset[i][1][\"answer\"]\n",
    "plot_qa_inference_results(img, question, answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46c36a34-8688-42ed-bf0c-7c6ecb364ab5",
   "metadata": {},
   "source": [
    "This sample has relation semantic type. We can see from the distribution of top predictions that the relation QA head is effective in anchoring to the text embeddings for left and right in the prompt."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecb38231-96fd-40bd-aadd-eb24c3e0b729",
   "metadata": {},
   "source": [
    "We can also construct our own questions (assuming the answer is present in the set of labels):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d3073f1-35a0-40af-b328-46898a82627c",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_qa_inference_results(img, \"What animal is inside the backpack?\", torch.tensor(dataset.answer2id[\"dog\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7526d12f-674b-488b-9c5a-178d8e0c37b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 1\n",
    "img = dataset[i][0]\n",
    "question = dataset[i][1][\"caption\"]\n",
    "answer = dataset[i][1][\"answer\"]\n",
    "plot_qa_inference_results(img, question, answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc6afc48-d0b5-42c7-8860-dcaaf60e73da",
   "metadata": {},
   "source": [
    "Note that in this case the detector associates the answer with the question's interrogative word. This phenomenon is also discussed in Fig. 5 of the [MDETR paper](https://arxiv.org/pdf/2104.12763.pdf)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81db236a-cce6-43aa-b08e-adc16857790d",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_qa_inference_results(img, \"What color is the man's bag?\", torch.tensor(dataset.answer2id[\"blue\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9911ef6-552b-42f2-93c0-25d7097aeca7",
   "metadata": {},
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db88b5dd-ac84-46fb-b5ff-648c42c224dd",
   "metadata": {},
   "source": [
    "In this notebook, we introduced MDETR and showed how to use the implementation from TorchMultimodal. We demonstrated how to perform inference on a couple samples for phrase grounding and VQA tasks. For instructions on how to perform fine-tuning for VQA using TorchMultimodal, you can also check out the README here (add link)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mm-backup-notebook",
   "language": "python",
   "name": "mm-backup-notebook"
  },
  "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}