{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import typing\n",
    "import io\n",
    "import os\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from urllib.request import urlretrieve\n",
    "\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "\n",
    "from models.modeling import VisionTransformer, CONFIGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(\"attention_data\", exist_ok=True)\n",
    "if not os.path.isfile(\"attention_data/ilsvrc2012_wordnet_lemmas.txt\"):\n",
    "    urlretrieve(\"https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt\", \"attention_data/ilsvrc2012_wordnet_lemmas.txt\")\n",
    "if not os.path.isfile(\"attention_data/ViT-B_16-224.npz\"):\n",
    "    urlretrieve(\"https://storage.googleapis.com/vit_models/imagenet21k+imagenet2012/ViT-B_16-224.npz\", \"attention_data/ViT-B_16-224.npz\")\n",
    "\n",
    "imagenet_labels = dict(enumerate(open('attention_data/ilsvrc2012_wordnet_lemmas.txt')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 224, 224])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test Image\n",
    "img_url = \"https://images.mypetlife.co.kr/content/uploads/2019/04/09192811/welsh-corgi-1581119_960_720.jpg\"\n",
    "urlretrieve(img_url, \"attention_data/img.jpg\")\n",
    "\n",
    "# Prepare Model\n",
    "config = CONFIGS[\"ViT-B_16\"]\n",
    "model = VisionTransformer(config, num_classes=1000, zero_head=False, img_size=224, vis=True)\n",
    "model.load_from(np.load(\"attention_data/ViT-B_16-224.npz\"))\n",
    "model.eval()\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((224, 224)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),\n",
    "])\n",
    "im = Image.open(\"attention_data/img.jpg\")\n",
    "x = transform(im)\n",
    "x.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "logits, att_mat = model(x.unsqueeze(0))\n",
    "\n",
    "att_mat = torch.stack(att_mat).squeeze(1)\n",
    "\n",
    "# Average the attention weights across all heads.\n",
    "att_mat = torch.mean(att_mat, dim=1)\n",
    "\n",
    "# To account for residual connections, we add an identity matrix to the\n",
    "# attention matrix and re-normalize the weights.\n",
    "residual_att = torch.eye(att_mat.size(1))\n",
    "aug_att_mat = att_mat + residual_att\n",
    "aug_att_mat = aug_att_mat / aug_att_mat.sum(dim=-1).unsqueeze(-1)\n",
    "\n",
    "# Recursively multiply the weight matrices\n",
    "joint_attentions = torch.zeros(aug_att_mat.size())\n",
    "joint_attentions[0] = aug_att_mat[0]\n",
    "\n",
    "for n in range(1, aug_att_mat.size(0)):\n",
    "    joint_attentions[n] = torch.matmul(aug_att_mat[n], joint_attentions[n-1])\n",
    "    \n",
    "# Attention from the output token to the input space.\n",
    "v = joint_attentions[-1]\n",
    "grid_size = int(np.sqrt(aug_att_mat.size(-1)))\n",
    "mask = v[0, 1:].reshape(grid_size, grid_size).detach().numpy()\n",
    "mask = cv2.resize(mask / mask.max(), im.size)[..., np.newaxis]\n",
    "result = (mask * im).astype(\"uint8\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prediction Label and Attention Map!\n",
      "\n",
      "0.49959 : Cardigan, Cardigan_Welsh_corgi\n",
      "0.47433 : Pembroke, Pembroke_Welsh_corgi\n",
      "0.01821 : dingo, warrigal, warragal, Canis_dingo\n",
      "0.00108 : basenji\n",
      "0.00077 : kelpie\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 559468 characters]\n",
      "text/plain": [
       "<Figure size 1152x1152 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(16, 16))\n",
    "\n",
    "ax1.set_title('Original')\n",
    "ax2.set_title('Attention Map')\n",
    "_ = ax1.imshow(im)\n",
    "_ = ax2.imshow(result)\n",
    "\n",
    "probs = torch.nn.Softmax(dim=-1)(logits)\n",
    "top5 = torch.argsort(probs, dim=-1, descending=True)\n",
    "print(\"Prediction Label and Attention Map!\\n\")\n",
    "for idx in top5[0, :5]:\n",
    "    print(f'{probs[0, idx.item()]:.5f} : {imagenet_labels[idx.item()]}', end='')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reference\n",
    "* [attention_flow](https://github.com/samiraabnar/attention_flow)\n",
    "* [vit-keras](https://github.com/faustomorales/vit-keras)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 540916 characters]\n",
      "text/plain": [
       "<Figure size 1152x1152 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 557544 characters]\n",
      "text/plain": [
       "<Figure size 1152x1152 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 557932 characters]\n",
      "text/plain": [
       "<Figure size 1152x1152 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 563900 characters]\n",
      "text/plain": [
       "<Figure size 1152x1152 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 562200 characters]\n",
      "text/plain": [
       "<Figure size 1152x1152 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 562704 characters]\n",
      "text/plain": [
       "<Figure size 1152x1152 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 564824 characters]\n",
      "text/plain": [
       "<Figure size 1152x1152 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 562576 characters]\n",
      "text/plain": [
       "<Figure size 1152x1152 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 562924 characters]\n",
      "text/plain": [
       "<Figure size 1152x1152 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 563660 characters]\n",
      "text/plain": [
       "<Figure size 1152x1152 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 561912 characters]\n",
      "text/plain": [
       "<Figure size 1152x1152 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 560104 characters]\n",
      "text/plain": [
       "<Figure size 1152x1152 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i, v in enumerate(joint_attentions):\n",
    "    # Attention from the output token to the input space.\n",
    "    mask = v[0, 1:].reshape(grid_size, grid_size).detach().numpy()\n",
    "    mask = cv2.resize(mask / mask.max(), im.size)[..., np.newaxis]\n",
    "    result = (mask * im).astype(\"uint8\")\n",
    "\n",
    "    fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(16, 16))\n",
    "    ax1.set_title('Original')\n",
    "    ax2.set_title('Attention Map_%d Layer' % (i+1))\n",
    "    _ = ax1.imshow(im)\n",
    "    _ = ax2.imshow(result)"
   ]
  }
 ],
 "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
