GradCAM.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Untitled1.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyOMWza48CXx7Hi4A4Utf1EB",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/SridharSola/Knowledge-Distillation-FER/blob/main/GradCAM.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TZrpcRjj5fwZ"
      },
      "outputs": [],
      "source": [
        "\"\"\"\n",
        "This method obtains the heatmap for gradients on the train set (or set given), superimposes the heatmap\n",
        "on the original image, and saves the superimposed image as a jpg file in a folder on the drive\n",
        "\"\"\"\n",
        "def GradCAM(net, train_loader):\n",
        "  net.cuda()\n",
        "  net.eval()\n",
        "\n",
        "  for j, (imgu, imgm, l, p1, p2) in enumerate(train_loader):\n",
        "      imgu = to_device(imgu, device)\n",
        "      pred1 = net(imgu)\n",
        "      indx = pred1.argmax(dim = 1)\n",
        "      # get the gradient of the output with respect to the parameters of the model\n",
        "      pred1[:, indx].backward()\n",
        "\n",
        "      # pull the gradients out of the model\n",
        "      gradients = net.get_activations_gradient()\n",
        "\n",
        "      # pool the gradients across the channels\n",
        "      pooled_gradients = torch.mean(gradients, dim=[0, 2, 3])\n",
        "\n",
        "      # get the activations of the last convolutional layer\n",
        "      activations = net.get_activations(imgu).detach()\n",
        "\n",
        "      # weight the channels by corresponding gradients\n",
        "      for i in range(512):\n",
        "          activations[:, i, :, :] *= pooled_gradients[i]\n",
        "          \n",
        "      # average the channels of the activations\n",
        "      heatmap = torch.mean(activations, dim=1).squeeze().cpu()\n",
        "      # relu on top of the heatmap\n",
        "      # expression (2) in https://arxiv.org/pdf/1610.02391.pdf\n",
        "      heatmap = np.maximum(heatmap, 0)\n",
        "      #Superimposing\n",
        "      img = cv2.imread(p1[0])\n",
        "      # normalize the heatmap\n",
        "      heatmap /= torch.max(heatmap)\n",
        "      heatmap = np.uint8(255 * heatmap)\n",
        "      heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]))\n",
        "      heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)\n",
        "      superimposed_img = heatmap * 0.4 + img\n",
        "      #cv2.imwrite('./map.jpg', superimposed_img)\n",
        "      if j%100:\n",
        "          cv2_imshow(superimposed_img)\n",
        "      # draw the heatmap\n",
        "      #plt.matshow(heatmap.squeeze())\n",
        "      cv2.imwrite('/content/drive/MyDrive/FER_Project/GradCam/' + str(j) + '.jpg', superimposed_img)"
      ]
    }
  ]
}