{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import torch.nn as nn\n",
    "import cv2\n",
    "from torchvision import models, transforms\n",
    "from timm.models.vision_transformer import vit_large_patch16_384\n",
    "from timm.models import create_model, resume_checkpoint, convert_splitbn_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = create_model(\n",
    "        'vit_large_patch16_384', \n",
    "        pretrained=True,\n",
    "        num_classes=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "block_weights = [] # we will save the conv layer weights in this list\n",
    "block_layers = [] # we will save the 49 conv layers in this list\n",
    "# get all the model children as list\n",
    "model_children = list(model.children())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_block = []  # put all sub block in the list\n",
    "for model_child in list(model.children()):\n",
    "    if type(model_child) != torch.nn.modules.container.ModuleList:\n",
    "        all_block.append(model_child)\n",
    "    else:\n",
    "        sub_children = list(model_child.children())\n",
    "        for sub_child in sub_children:\n",
    "            all_block.append(sub_child)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#len(all_block), all_block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 186744 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 512, 512])\n",
      "torch.Size([1, 3, 512, 512])\n"
     ]
    }
   ],
   "source": [
    "# read and visualize an image\n",
    "img = cv2.imread(\"images/dog.png\")\n",
    "img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "plt.imshow(img)\n",
    "plt.show()\n",
    "# define the transforms\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToPILImage(),\n",
    "    transforms.Resize((512,512)),  # reszie the image as (1024, 1024) or (2048, 2048) to obtain high resolution features if you have large CPU or GPU memory\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "img = np.array(img)\n",
    "# apply the transforms\n",
    "img = transform(img)\n",
    "print(img.size())\n",
    "# unsqueeze to add a batch dimension\n",
    "img = img.unsqueeze(0)\n",
    "print(img.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pass the image through all the layers\n",
    "results = []\n",
    "results = [all_block[0](img)]\n",
    "for i in range(1, len(all_block)-2):\n",
    "    # pass the result from the last layer to the next layer\n",
    "    #print(i)\n",
    "    results.append(all_block[i](results[-1]))\n",
    "# make a copy of the `results`\n",
    "outputs = results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#outputs[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving layer 0 feature maps...\n",
      "Saving layer 1 feature maps...\n",
      "Saving layer 2 feature maps...\n",
      "Saving layer 3 feature maps...\n",
      "Saving layer 4 feature maps...\n"
     ]
    }
   ],
   "source": [
    "# visualize 36 features from each layer \n",
    "# (although there are more feature maps in the upper layers)\n",
    "all_layers_num = 5   # change the number of length to control how many layers you want to visualize\n",
    "for num_layer in range(all_layers_num): \n",
    "    plt.figure(figsize=(30, 30))\n",
    "    layer_viz = outputs[num_layer]#[0, :, :, :]\n",
    "    layer_viz = layer_viz.data\n",
    "    layer_viz = layer_viz.squeeze(0).transpose(0, 1)\n",
    "    layer_viz = layer_viz.reshape(1024, 32, 32)  # chage the reshape here\n",
    "    #print(layer_viz.size())\n",
    "    for i, filter in enumerate(layer_viz):\n",
    "        if i == 36: # we will visualize only 8x8 blocks from each layer\n",
    "            break\n",
    "        plt.subplot(6, 6, i + 1)\n",
    "        plt.imshow(filter, cmap='gray')\n",
    "        plt.axis(\"off\")\n",
    "    print(f\"Saving layer {num_layer} feature maps...\")\n",
    "    plt.savefig(f\"data/outputs/vit/layer_{num_layer}.png\")  # change the path to save the feature maps\n",
    "    #plt.show() # use this line to show the figure in jupyter notebook\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
