visualize_attention_map.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "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 utils.dataset import *\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "\n",
    "from models.modeling import VisionTransformer, CONFIGS\n",
    "from torch.utils.data import DataLoader, RandomSampler, DistributedSampler, SequentialSampler\n",
    "import torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test Image\n",
    "\n",
    "# Prepare Model\n",
    "config = CONFIGS[\"ViT-B_16\"]\n",
    "model = VisionTransformer(config, num_classes=1938, zero_head=False, img_size=448, vis=True)\n",
    "checkpoint=torch.load(\"/home/deep/junwang/MTGCV_V/soygbl_base_linear_ns290400_2e-2_bs4_cos.log_checkpoint.bin\")\n",
    "model.load_state_dict(checkpoint)\n",
    "model.eval()\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((600, 600)),\n",
    "    transforms.CenterCrop((448,448)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = CONFIGS[\"ViT-B_16\"]\n",
    "model = make_model(config, None, zero_head=True, num_classes=200, vis=True)\n",
    "checkpoint=torch.load('/home/ubuntu/junwang/paper/aaai2021/inter_vit_jun/output/soybean200_InterViT_checkpoint.bin')\n",
    "model.load_state_dict(checkpoint)\n",
    "model.eval()\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((600, 600)),\n",
    "    transforms.CenterCrop((448,448)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_name=\"soybean2000\"\n",
    "dataset=eval(data_name)(root='./data/'+data_name, is_train=False, transform=transform)\n",
    "train_sampler = RandomSampler(dataset)\n",
    "test_loader = DataLoader(dataset,sampler=train_sampler,batch_size=1,num_workers=4,pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_colors(num_colors):\n",
    "    \"\"\"\n",
    "    Generate distinct value by sampling on hls domain.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    num_colors: int\n",
    "        Number of colors to generate.\n",
    "\n",
    "    Returns\n",
    "    ----------\n",
    "    colors_np: np.array, [num_colors, 3]\n",
    "        Numpy array with rows representing the colors.\n",
    "\n",
    "    \"\"\"\n",
    "    colors=[]\n",
    "    for i in np.arange(0., 360., 360. / num_colors):\n",
    "        hue = i/360.\n",
    "        lightness = 0.5\n",
    "        saturation = 0.9\n",
    "        colors.append(colorsys.hls_to_rgb(hue, lightness, saturation))\n",
    "    colors_np = np.array(colors)*255.\n",
    "\n",
    "    return colors_np\n",
    "\n",
    "def show_att_on_image(img, mask, output):\n",
    "    \"\"\"\n",
    "    Convert the grayscale attention into heatmap on the image, and save the visualization.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    img: np.array, [H, W, 3]\n",
    "        Original colored image.\n",
    "    mask: np.array, [H, W]\n",
    "        Attention map normalized by subtracting min and dividing by max.\n",
    "    output: str\n",
    "        Destination image (path) to save.\n",
    "\n",
    "    Returns\n",
    "    ----------\n",
    "    Save the result to output.\n",
    "c\n",
    "    \"\"\"\n",
    "    img_h, img_w = img.size[0], img.size[1]\n",
    "    plt.subplots(nrows=1, ncols=1, figsize=(0.02*img_h, 0.02*img_w))\n",
    "    plt.axis('off')\n",
    "    plt.imshow(img, alpha=1)\n",
    "    normed_mask = mask / mask.max()\n",
    "    normed_mask = (normed_mask * 255).astype('uint8')\n",
    "    plt.imshow(normed_mask, alpha=0.5, interpolation='nearest', cmap='jet')\n",
    "    \n",
    "    plt.savefig(output)\n",
    "    plt.close()\n",
    "    #heatmap = cv2.applyColorMap(np.uint8(255*mask), cv2.COLORMAP_JET)\n",
    "    #heatmap = np.float32(heatmap) / 255\n",
    "\n",
    "    # add heatmap onto the image\n",
    "    #merged = heatmap + np.float32(img)\n",
    "\n",
    "    # re-scale the image\n",
    "    #merged = merged / np.max(merged)\n",
    "    #cv2.imwrite(output, np.uint8(255 * merged))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig_rows = 1\n",
    "fig_cols = 1\n",
    "f_assign, axarr_assign = plt.subplots(fig_rows, fig_cols, figsize=(fig_cols*2,fig_rows*2))\n",
    "\n",
    "root=os.path.join('./visualization',data_name)\n",
    "with torch.no_grad():\n",
    "    for i,data in enumerate(test_loader):\n",
    "        if i>=500:break\n",
    "        dir_path_attn=os.path.join(root,str(i),'attention_map')\n",
    "        os.makedirs(dir_path_attn, exist_ok=True)\n",
    "        x , label=data[0].squeeze(0),data[1]\n",
    "        logits,  att_mat= model(x.unsqueeze(0))\n",
    "        att_mat = torch.stack(att_mat).squeeze(1)\n",
    "        ### visualize attention map\n",
    "        att_mat = torch.mean(att_mat, dim=1)\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",
    "        \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(), (448,448))[..., np.newaxis]\n",
    "        mask = mask[:,:,0]\n",
    "        \n",
    "        save_input = transforms.Normalize(mean=(0, 0, 0),std=(1/0.229, 1/0.224, 1/0.225))(x.data.cpu())\n",
    "        save_input = transforms.Normalize(mean=(-0.485, -0.456, -0.406),std=(1, 1, 1))(save_input)\n",
    "        save_input = torch.nn.functional.interpolate(save_input.unsqueeze(0), size=(448, 448), mode='bilinear', align_corners=False).squeeze(0)\n",
    "        img = torchvision.transforms.ToPILImage()(save_input)\n",
    "        \n",
    "        show_att_on_image(img, mask, os.path.join(dir_path_attn,'heatmap.png'))\n",
    "        \n",
    "\n",
    "        img.save(dir_path_attn+\"/input.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "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 (ipykernel)",
   "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}