semantic_encoding/_utils/tsne_attention_plot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"..\")\n",
    "\n",
    "import os\n",
    "import argparse\n",
    "import multiprocessing\n",
    "import pdb\n",
    "import random\n",
    "from datetime import datetime\n",
    "\n",
    "from PIL import Image\n",
    "from coco_captions_dataset import CocoCaptionsDataset\n",
    "from utils import load_checkpoint, convert_sentences_to_word_idxs\n",
    "\n",
    "import numpy as np\n",
    "from sklearn.manifold import TSNE\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.patches import ArrowStyle\n",
    "from matplotlib.patches import FancyBboxPatch\n",
    "import seaborn as sns\n",
    "sns.set(\n",
    "    style=\"darkgrid\"\n",
    ")\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch\n",
    "import torchvision.transforms as T\n",
    "import torch.backends.cudnn as cudnn\n",
    "\n",
    "IMAGENET_MEAN   = [0.485, 0.456, 0.406]\n",
    "IMAGENET_STD    = [0.229, 0.224, 0.225]\n",
    "FEAT_EMBED_SIZE = 128\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "cudnn.benchmark = True  # set to true only if inputs to model are fixed size; otherwise lot of computational overhead"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_checkpoint_filename = '../models/vdan_pretrained_model.pth'\n",
    "\n",
    "image_path = '000000401862.jpg'\n",
    "image_long_path = '../resources/COCO_2017/val2017/' + image_path\n",
    "batch_size = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "captions_val_fname = '../resources/COCO_2017/annotations/captions_val2017.json'\n",
    "val_data_path = '../resources/COCO_2017/val2017/'\n",
    "num_workers = int(multiprocessing.cpu_count()) # Using 80% of CPU cores\n",
    "print('Using %d CPU cores...' % num_workers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('[{}] Loading saved model weights: {}...'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), model_checkpoint_filename))\n",
    "_, model, optimizer_state_dict, word_map, model_params, train_params = load_checkpoint(model_checkpoint_filename)\n",
    "model.to(device)\n",
    "model.eval()\n",
    "print('[{}] Done!\\n'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('[{}] Computing image features for {} ...'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), image_long_path))\n",
    "img_transform = T.Compose([T.Resize((224,224)),\n",
    "                           T.ToTensor(),\n",
    "                           T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)])\n",
    "\n",
    "img_original = Image.open(image_long_path).convert('RGB')\n",
    "img = img_transform(img_original).unsqueeze(0).to(device)\n",
    "\n",
    "# Extracting image features\n",
    "with torch.no_grad():\n",
    "    imgs_feats, resnet_output = model.get_img_embedding(img)\n",
    "resnet_output = torch.cat(batch_size*[resnet_output])\n",
    "\n",
    "print('[{}] Done!\\n'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = CocoCaptionsDataset(root = val_data_path,\n",
    "                                annFile = captions_val_fname,\n",
    "                                word_map = word_map,\n",
    "                                num_sentences=10,\n",
    "                                img_transform=img_transform,\n",
    "                                annotations_transform=T.ToTensor(),\n",
    "                                dataset_proportion=1.,\n",
    "                                generate_negatives=False)\n",
    "\n",
    "test_dataloader = torch.utils.data.DataLoader(test_data,\n",
    "                                            batch_size=batch_size,\n",
    "                                            num_workers=num_workers,\n",
    "                                            shuffle=True)\n",
    "\n",
    "print('[{}] Extracting COCO annotations features...'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S')))\n",
    "\n",
    "num_imgs = len(test_data)\n",
    "docs_feats = np.ndarray((num_imgs,FEAT_EMBED_SIZE), dtype=np.float32)\n",
    "docs_attention = np.ndarray((num_imgs,test_data.num_sentences), dtype=np.float32)\n",
    "docs_text = []\n",
    "conditioning_image_id = None\n",
    "\n",
    "for i, (img_path, captions, _, documents, sentences_per_document, words_per_sentence, is_negative_sample) in enumerate(tqdm(test_dataloader)):\n",
    "    docs_text += [list(x) for x in np.array(captions).T]\n",
    "    if image_path in img_path:\n",
    "        conditioning_image_id = i*batch_size + list(img_path).index(image_path)\n",
    "        print(\"Found image! {} ID:{} Docs:\\n{}\".format(image_path, conditioning_image_id, docs_text[conditioning_image_id]))\n",
    "        \n",
    "    documents = documents.squeeze(1).to(device)  # (batch_size, sentence_limit, word_limit)\n",
    "    sentences_per_document = sentences_per_document.to(device) # (batch_size)\n",
    "    words_per_sentence = words_per_sentence.to(device)  # (batch_size, sentence_limit)\n",
    "    \n",
    "    # Extracting COCO annotations features\n",
    "    with torch.no_grad():\n",
    "        annotation_feats, _, sentence_alphas = model.get_text_embedding(documents, sentences_per_document, words_per_sentence, resnet_output)\n",
    "    docs_feats[i*batch_size:min((i+1)*batch_size,num_imgs),:] = annotation_feats.detach().cpu().numpy()\n",
    "    docs_attention[i*batch_size:min((i+1)*batch_size,num_imgs),:] = sentence_alphas.detach().cpu().numpy()\n",
    "    \n",
    "dots = np.matmul(docs_feats, imgs_feats.detach().cpu().numpy().transpose()).squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tsne = TSNE(perplexity=30, n_iter=5000)\n",
    "docs_embedded_original = tsne.fit_transform(docs_feats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "zoom_offset = np.array([-5, -7])\n",
    "zoom_center = np.average(docs_embedded_original, axis=0, weights=dots.clip(0, 1))\n",
    "zoom_radius = 20\n",
    "\n",
    "x1 = zoom_center[0]/np.linalg.norm(zoom_center)\n",
    "x2 = 1\n",
    "y1 = zoom_center[1]/np.linalg.norm(zoom_center)\n",
    "y2 = 0\n",
    "rot = np.array(\n",
    "    [[  x1*x2 + y1*y2 , x1*y2 - x2*y1],\n",
    "     [-(x1*y2 - x2*y1), x1*x2 + y1*y2]])\n",
    "print(zoom_center)\n",
    "\n",
    "rot = np.array([[1.0,0.0], [0.0,1.0]])\n",
    "\n",
    "zoom_center = (rot @ zoom_center.T) + zoom_offset\n",
    "print(zoom_center)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "docs_embedded = (rot @docs_embedded_original.T).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams[\"figure.figsize\"] = (12,6)\n",
    "plt.rcParams[\"figure.dpi\"] = 300\n",
    "\n",
    "plot_ax = plt.axes([0.05,0.1, 0.35, 0.7], frameon=True)\n",
    "scatter = sns.scatterplot(docs_embedded[:,0], docs_embedded[:,1], hue=dots, s=15, linewidth=0, palette=\"RdBu\", legend=\"brief\", ax=plot_ax)\n",
    "\n",
    "for t in scatter.get_legend().texts:\n",
    "    if float(t.get_text()) > 1:\n",
    "        t.set_text(' 1.0')\n",
    "    if float(t.get_text()) < -1:\n",
    "        t.set_text('-1.0')\n",
    "    if t.get_text()[0] not in [' ', '-']:\n",
    "        t.set_text(' '+t.get_text())\n",
    "\n",
    "text_ax = plt.axes([0.45,0.25, 0.25, 0.5], frameon=True)\n",
    "data_to_axis = (text_ax.transAxes + text_ax.transData).inverted()\n",
    "scatter_zoom = sns.scatterplot(docs_embedded[:,0], docs_embedded[:,1], hue=dots, palette=\"RdBu\", legend=False, ax=text_ax)\n",
    "\n",
    "x1, x2 = zoom_center[0]-zoom_radius, zoom_center[0]+zoom_radius\n",
    "y1, y2 = zoom_center[1]-zoom_radius, zoom_center[1]+zoom_radius\n",
    "\n",
    "selected_points = []\n",
    "for i in range(docs_embedded.shape[0]):\n",
    "    if i == conditioning_image_id:\n",
    "        continue\n",
    "    x = docs_embedded[i,0]\n",
    "    y = docs_embedded[i,1]\n",
    "    if x > x1 and x < x2 and y > y1 and y < y2 and dots[i] > 0.9:\n",
    "        selected_points += [{'x':x, 'y':y ,'dot': dots[i]\n",
    "                             ,'atts':  docs_attention[i]\n",
    "                            #,'text': '\\n'.join(docs_text[i])\n",
    "                             ,'texts': [\"{:.2f} | {}\".format(docs_attention[i][j], docs_text[i][j]) for j in range(test_data.num_sentences)]\n",
    "                            }]\n",
    "n_texts = 3\n",
    "selected_points = random.choices(selected_points, k=n_texts-1)\n",
    "selected_points.sort(reverse=True, key=lambda x: x['dot']*60 + (x['x']-zoom_center[0]) + x['y']-zoom_center[1])\n",
    "selected_points = [{\"x\": docs_embedded[conditioning_image_id,0]\n",
    "                   ,\"y\": docs_embedded[conditioning_image_id,1]\n",
    "                   ,\"dot\": dots[conditioning_image_id]\n",
    "                   ,'atts': docs_attention[conditioning_image_id]\n",
    "                  #,\"text\": \"\\n\".join(map(lambda x: x.rstrip(), docs_text[conditioning_image_id]))\n",
    "                   ,'texts':  [\"{:.2f} | {}\".format(docs_attention[conditioning_image_id][j], docs_text[conditioning_image_id][j]) for j in range(test_data.num_sentences)]\n",
    "                   }] + selected_points\n",
    "                  \n",
    "\n",
    "text_y_offset = (zoom_radius*5)/n_texts\n",
    "text_x = zoom_center[0] + zoom_radius*1.2\n",
    "text_y = zoom_center[1] + zoom_radius - text_y_offset/5\n",
    "text_line_offset_y = (text_y_offset/test_data.num_sentences) * 0.76#* 0.795\n",
    "text_box_offset_y0 = 4\n",
    "text_box_offset_y1 = 2\n",
    "arrowstyle = ArrowStyle('fancy', head_width=1.0, head_length=1.0, tail_width=0.8)\n",
    "\n",
    "for i in reversed(range(n_texts)):\n",
    "    center_offset = float(n_texts-i)/n_texts * text_box_offset_y0 + float(i)/n_texts * text_box_offset_y1\n",
    "    scatter_zoom.annotate(\"\"\n",
    "                         ,(selected_points[i]['x'],selected_points[i]['y'])\n",
    "                         ,(text_x, text_y-i*text_y_offset + text_y_offset*0.3125)\n",
    "                         ,arrowprops = {'arrowstyle': arrowstyle\n",
    "                                       ,'color': '0.95' if i != 0 else 'b'\n",
    "                                       ,'ec': 'b' if i != 0 else 'white'\n",
    "                                       ,'connectionstyle': \"arc3,rad={}\".format(float(n_texts-i)/n_texts * 0.3 + float(i)/n_texts * -0.3)})\n",
    "    \n",
    "    scatter_zoom.text(text_x - 0.5\n",
    "                     ,text_y-i*text_y_offset+ text_y_offset*0.8075\n",
    "                     ,\"Dot: {:0.3f}\".format(selected_points[i]['dot'])\n",
    "                     ,horizontalalignment='left', size='small', color='black', weight='semibold')\n",
    "    \n",
    "    bgbox = scatter_zoom.annotate(\"\\n\".join(selected_points[i]['texts'])\n",
    "                         ,(selected_points[i]['x'],selected_points[i]['y'])\n",
    "                         ,(text_x, text_y - i*text_y_offset)\n",
    "                         , horizontalalignment='left', size='small'\n",
    "                         , bbox=dict(boxstyle='round', pad=0.7, fc='0.95' if i != 0 else 'b')\n",
    "                         , weight='normal' if i != 0 else 'semibold'\n",
    "                         , linespacing=1.5\n",
    "                         , alpha=0)\n",
    "\n",
    "    a_min = selected_points[i]['atts'].min()\n",
    "    a_max = selected_points[i]['atts'].max()\n",
    "    for j in range(test_data.num_sentences):\n",
    "        scatter_zoom.annotate(selected_points[i]['texts'][j]\n",
    "                             ,(selected_points[i]['x'], selected_points[i]['y'])\n",
    "                             ,(text_x, text_y - i*text_y_offset + (test_data.num_sentences-j-1)*text_line_offset_y)\n",
    "                             , horizontalalignment='left', size='small'\n",
    "                             , weight='normal' if i != 0 else 'semibold'\n",
    "                             , color='black' if i != 0 else 'white'\n",
    "                             , alpha=((selected_points[i]['atts'][j]-a_min)/(a_max-a_min) + 0.4)/ 1.4)\n",
    "text_ax.set_xlim(x1, x2)\n",
    "text_ax.set_ylim(y1, y2)\n",
    "text_ax.set_xticklabels('')\n",
    "text_ax.set_yticklabels('')\n",
    "\n",
    "plot_ax.indicate_inset_zoom(text_ax)\n",
    "\n",
    "img_ax = plt.axes([0.004, 0.82, 0.175, 0.175], frameon=True)\n",
    "img_ax.imshow(img_original)\n",
    "img_ax.axis('off')\n",
    "\n",
    "plt.suptitle('T-SNE', x=0.23, y=0.925, size=28, weight='light')\n",
    "\n",
    "img_name = os.path.splitext(os.path.basename(image_path))[0]\n",
    "plt.savefig('tsne_att_' + img_name + '.png', dpi=400, bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('[{}] Done!\\n'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S')))"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}