visualize_group_features.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import typing\n",
    "import io\n",
    "import os\n",
    "import colorsys\n",
    "import argparse\n",
    "import torchvision\n",
    "import numpy as np\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n",
    "from urllib.request import urlretrieve\n",
    "from utils.dataset import *\n",
    "\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "from models.make_model import make_model\n",
    "\n",
    "from models.modeling import VisionTransformer, CONFIGS\n",
    "from torch.utils.data import DataLoader, RandomSampler, DistributedSampler, SequentialSampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "===========building transformer===========\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# Prepare Model\n",
    "config = CONFIGS[\"R50-ViT-B_16\"]\n",
    "model = make_model(config, None, zero_head=True, num_classes=4, vis=True)\n",
    "checkpoint=torch.load('./trained_models/AFD_tp02_cw4_rp_sd3567_6500_5_10_20_5e-3_bs16_wd5e-5_ly4_np5.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.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_name=\"AFD\"\n",
    "dataset=eval(data_name)(root='./data/'+data_name, is_train=True, transform=transform)\n",
    "train_sampler = RandomSampler(dataset)\n",
    "train_loader = DataLoader(dataset,sampler=train_sampler,batch_size=1,num_workers=4,pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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",
    "\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))\n",
    "\n",
    "def plot_assignment(root, assign_hard, num_parts):\n",
    "    \"\"\"\n",
    "    Blend the original image and the colored assignment maps.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    root: str\n",
    "        Root path for saving visualization results.\n",
    "    assign_hard: np.array, [H, W]\n",
    "        Hard assignment map (int) denoting the deterministic assignment of each pixel. Generated via argmax.\n",
    "    num_parts: int, number of object parts.\n",
    "\n",
    "    Returns\n",
    "    ----------\n",
    "    Save the result to root/assignment.png.\n",
    "\n",
    "    \"\"\"\n",
    "    # generate the numpy array for colors\n",
    "    colors = generate_colors(num_parts)\n",
    "\n",
    "    # coefficient for blending\n",
    "    coeff = 0.4\n",
    "\n",
    "    # load the input as RGB image, convert into numpy array\n",
    "    input = Image.open(os.path.join(root, 'input.png')).convert('RGB')\n",
    "    input_np = np.array(input).astype(float)\n",
    "\n",
    "    # blending by each pixel\n",
    "    for i in range(assign_hard.shape[0]):\n",
    "        for j in range(assign_hard.shape[1]):\n",
    "            assign_ij = assign_hard[i][j]\n",
    "            input_np[i, j] = (1-coeff) * input_np[i, j] + coeff * colors[assign_ij]\n",
    "\n",
    "    # save the resulting image\n",
    "    im = Image.fromarray(np.uint8(input_np))\n",
    "    im.save(os.path.join(root, 'assignment.png'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAJ0AAACQCAYAAAASuGkIAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8/fFQqAAAACXBIWXMAAAsTAAALEwEAmpwYAAAHEElEQVR4nO3dX4hcZxnH8e/PJgWNfyJu1FqNVIiNEVpJ1xrBPxFR070JhV60FYNBWJQqXupVe9ErLwQptQ1LCaE37Y2lVknrneaiRrorSZq2KNuWxrWFJFYitaJs+3hxzsq4md09OfvOM9k3vw8MzMz5M8/J+TGzk/c88yoiMMv0jnEXYFceh87SOXSWzqGzdA6dpXPoLN2aoZN0WNJZSadXWC5J90mal3RK0u7yZVpNurzTHQH2rbL8FmBHe5sGHlx/WVazNUMXEceA11dZZT/wcDSOA1slXVOqQKtPib/prgX+MvB4oX3ObKhNBfahIc8NHVuTNE3zEcyWLVtu2rlzZ4GXt3GYm5s7HxHb+mxbInQLwMcGHn8UeHXYihExA8wATE5OxuzsbIGXt3GQ9ErfbUt8vD4BHGi/xe4BLkTEawX2a5Va851O0iPAXmBC0gJwD7AZICIOAUeBKWAeeBM4OKpirQ5rhi4i7lhjeQB3FavIqucRCUvn0Fk6h87SOXSWzqGzdA6dpXPoLJ1DZ+kcOkvn0Fk6h87SOXSWzqGzdJ1CJ2mfpD+1HV8/HrL8fZJ+JemkpOck+fImW1GXFsSrgJ/TdH3tAu6QtGvZancBz0fEjTTX3v1U0tWFa7VKdHmnuxmYj4iXIuI/wKM0HWCDAniPJAHvpukeWyxaqVWjS+i6dHvdD3yKpjfiWeCHEfF2kQqtOl1C16Xb6xvACeAjwGeA+yW996IdSdOSZiXNnjt37hJLtVp0CV2Xbq+DwGNtw/U88DJwUX9hRMxExGRETG7b1qt7zSrQJXTPADskXdd+ObidpgNs0BngqwCSPgRcD7xUslCrR5fGnEVJ3wd+A1wFHI6I5yR9t11+CLgXOCLpWZqP4x9FxPkR1m0bWKdm64g4StNqOPjcoYH7rwJfL1ua1cojEpbOobN0Dp2lc+gsnUNn6Rw6S+fQWTqHztI5dJbOobN0Dp2lc+gsnUNn6Yp0g7Xr7JV0ou0G+13ZMq0mXX5dfakb7Gs0VxE/I+mJiHh+YJ2twAPAvog4I+mDI6rXKlCqG+xOmsvVzwBExNmyZVpNSnWDfRJ4v6TfSpqTdKBUgVafLlcOd+kG2wTcRNMn8U7g95KOR8Sf/29HA3ODbd++/dKrtSqU6gZbAJ6KiH+2vRHHgBuX78jdYAblusF+CXxR0iZJ7wI+B7xQtlSrRZFusIh4QdJTwCngbeChiBg6/bqZmqm98nnqzY1N0lxETPbZ1iMSls6hs3QOnaVz6CydQ2fpHDpL59BZOofO0jl0ls6hs3QOnaVz6Cxdscacdr3PSnpL0m3lSrTalJqmaWm9n9BcAmW2olKNOQA/AH4BuCnHVlWkMUfStcCtwCHM1lBqmqaf0cwd8daqO/I0TUa3brAujTmTwKPNJIhMAFOSFiPi8cGVImIGmIHmyuGeNdsG1yV0/2vMAf5K05hz5+AKEXHd0n1JR4BfLw+c2ZJS0zSZdVZkmqZlz397/WVZzTwiYekcOkvn0Fk6h87SOXSWzqGzdA6dpXPoLJ1DZ+kcOkvn0Fk6h87SOXSWrkg3mKRvSjrV3p6WdNEvq5stKdUN9jLw5Yi4AbiX9upgs2GKdINFxNMR8ff24XGaS9rNhio1TdOg7wBPrqcoq1upaZqaFaWv0ITuCyss9zRNVmyaJiTdADwE7I+Ivw3bkadpMig0TZOk7cBjwLeWT0JntlypbrC7gQ8AD7S9r4t9Z1Ox+nmaJuvF0zTZhuLQWTqHztI5dJbOobN0Dp2lc+gsnUNn6Rw6S+fQWTqHztI5dJbOobN0pbrBJOm+dvkpSbvLl2q1KNUNdguwo71NAw8WrtMqUmpusP3Aw9E4DmyVdE3hWq0SpbrBLrVjzK5gpbrBOnWMDXaDAf+WdLrD629UE8D5cRcxQtf33bDU3GCdOsYG5waTNFtzH8WVcHx9ty3SDdY+PtB+i90DXIiI1/oWZXUr1Q12FJgC5oE3gYOjK9k2urF1g0mabj9uq+TjW2XbcYXOrlweBrN0Iw9d7UNoHY5vr6QLkk60t7vHUWcfkg5LOrvSf231PncRMbIbzRePF4FPAFcDJ4Fdy9aZovlpMQF7gD+MsqYxHN9empm+x15vj+P7ErAbOL3C8l7nbtTvdLUPoXU5vg0rIo4Br6+ySq9zN+rQ1T6E1rX2z0s6KelJSZ/OKS1Fr3PXaTr1dSg2hHaZ6lL7H4GPR8QbkqaAx2muxqlBr3M36ne6YkNol6k1a4+If0TEG+39o8BmSRN5JY5Ur3M36tDVPoTW5QcjP6z2R/sk3Uzzbz70l0o3oF7nbqQfr1H5EFrH47sN+J6kReBfwO3RfvW73El6hObb94SkBeAeYDOs79x5RMLSeUTC0jl0ls6hs3QOnaVz6CydQ2fpHDpL59BZuv8CLhhbE+FsPYMAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 144x144 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "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(train_loader):\n",
    "        if i>=200:break\n",
    "        dir_path_assign=os.path.join(root,str(i),'assignments')\n",
    "        dir_path_attn=os.path.join(root,str(i),'attention_map')\n",
    "        os.makedirs(dir_path_assign, exist_ok=True)\n",
    "        os.makedirs(dir_path_attn, exist_ok=True)\n",
    "        x , label=data[0].squeeze(0),data[1]\n",
    "        logits, assign, 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",
    "        size=assign.shape\n",
    "        assign_attn=torch.mul(v[0, 1:].unsqueeze(1),assign.flatten(-2))\n",
    "        assign_attn=assign_attn.view(*size)\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",
    "        img.save(dir_path_assign+\"/input.png\")\n",
    "        img.save(dir_path_attn+\"/input.png\")\n",
    "        assign_reshaped = torch.nn.functional.interpolate(assign.data.cpu(), size=(448, 448), mode='bilinear', align_corners=False)\n",
    "        assign_attn_reshaped = torch.nn.functional.interpolate(assign_attn.data.cpu(), size=(448, 448), mode='bilinear', align_corners=False)\n",
    "        _, assign = torch.max(assign_reshaped, 1)\n",
    "        _, assign_attn = torch.max(assign_attn_reshaped, 1)\n",
    "        attn_mask = torch.sum(assign_attn_reshaped, 1)\n",
    "        # colorize and save the assignment\n",
    "        plot_assignment(dir_path_assign, assign.squeeze(0).numpy(), 5)\n",
    "        show_att_on_image(img, attn_mask.squeeze(0).numpy(), os.path.join(dir_path_attn,'heatmap.png'))\n",
    "\n",
    "        # plot the assignment for each dictionary vector\n",
    "        for i in range(5):\n",
    "            img = torch.nn.functional.interpolate(assign_reshaped.data[:, i].cpu().unsqueeze(0), size=(448, 448), mode='bilinear', align_corners=False)\n",
    "            img = torchvision.transforms.ToPILImage()(img.squeeze(0))\n",
    "            img.save(os.path.join(dir_path_assign, 'part_'+str(i)+'.png'))\n",
    "            \n",
    "        for i in range(5):\n",
    "            img = torch.nn.functional.interpolate(assign_attn_reshaped.data[:, i].cpu().unsqueeze(0), size=(448, 448), mode='bilinear', align_corners=False)\n",
    "            img = torchvision.transforms.ToPILImage()(img.squeeze(0))\n",
    "            img.save(os.path.join(dir_path_attn, 'part_'+str(i)+'.png'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.is_available()"
   ]
  }
 ],
 "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}