{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "09a56592-43a1-4f58-97ce-5fbf81ebcc53",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: grad-cam in /opt/conda/lib/python3.10/site-packages (1.4.8)\n",
      "Requirement already satisfied: opencv-python in /opt/conda/lib/python3.10/site-packages (from grad-cam) (4.8.0.74)\n",
      "Requirement already satisfied: Pillow in /opt/conda/lib/python3.10/site-packages (from grad-cam) (9.4.0)\n",
      "Requirement already satisfied: ttach in /opt/conda/lib/python3.10/site-packages (from grad-cam) (0.0.3)\n",
      "Requirement already satisfied: numpy in /opt/conda/lib/python3.10/site-packages (from grad-cam) (1.23.5)\n",
      "Requirement already satisfied: torchvision>=0.8.2 in /opt/conda/lib/python3.10/site-packages (from grad-cam) (0.14.1+cu116)\n",
      "Requirement already satisfied: torch>=1.7.1 in /opt/conda/lib/python3.10/site-packages (from grad-cam) (1.13.1+cu116)\n",
      "Requirement already satisfied: matplotlib in /opt/conda/lib/python3.10/site-packages (from grad-cam) (3.7.1)\n",
      "Requirement already satisfied: scikit-learn in /opt/conda/lib/python3.10/site-packages (from grad-cam) (1.2.2)\n",
      "Requirement already satisfied: tqdm in /opt/conda/lib/python3.10/site-packages (from grad-cam) (4.65.0)\n",
      "Requirement already satisfied: typing-extensions in /opt/conda/lib/python3.10/site-packages (from torch>=1.7.1->grad-cam) (4.5.0)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.10/site-packages (from torchvision>=0.8.2->grad-cam) (2.28.2)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /opt/conda/lib/python3.10/site-packages (from matplotlib->grad-cam) (2.8.2)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib->grad-cam) (23.0)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib->grad-cam) (4.39.3)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib->grad-cam) (1.4.4)\n",
      "Requirement already satisfied: pyparsing>=2.3.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib->grad-cam) (3.0.9)\n",
      "Requirement already satisfied: contourpy>=1.0.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib->grad-cam) (1.0.7)\n",
      "Requirement already satisfied: cycler>=0.10 in /opt/conda/lib/python3.10/site-packages (from matplotlib->grad-cam) (0.11.0)\n",
      "Requirement already satisfied: scipy>=1.3.2 in /opt/conda/lib/python3.10/site-packages (from scikit-learn->grad-cam) (1.10.1)\n",
      "Requirement already satisfied: joblib>=1.1.1 in /opt/conda/lib/python3.10/site-packages (from scikit-learn->grad-cam) (1.2.0)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /opt/conda/lib/python3.10/site-packages (from scikit-learn->grad-cam) (3.1.0)\n",
      "Requirement already satisfied: six>=1.5 in /opt/conda/lib/python3.10/site-packages (from python-dateutil>=2.7->matplotlib->grad-cam) (1.16.0)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision>=0.8.2->grad-cam) (1.26.15)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision>=0.8.2->grad-cam) (3.4)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision>=0.8.2->grad-cam) (3.1.0)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision>=0.8.2->grad-cam) (2022.12.7)\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install grad-cam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "86bd3d27-36f9-417d-a6e6-ee619d6dca14",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "warnings.simplefilter('ignore')\n",
    "import math\n",
    "from torchvision import datasets, models, transforms\n",
    "import torch\n",
    "import torch.functional as F\n",
    "import numpy as np\n",
    "import requests\n",
    "import torchvision\n",
    "from torch.utils.data import DataLoader\n",
    "from torch import nn\n",
    "from time import time\n",
    "from PIL import Image\n",
    "from tqdm.notebook import tqdm\n",
    "from pytorch_grad_cam.utils.image import show_cam_on_image, preprocess_image\n",
    "from pytorch_grad_cam import GradCAM, HiResCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM, EigenGradCAM, FullGrad, RandomCAM\n",
    "from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget\n",
    "from pytorch_grad_cam.utils.image import show_cam_on_image\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "012bf3c2-0550-444e-8c4e-4448f6b4cda0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11\n"
     ]
    }
   ],
   "source": [
    "data_transform = transforms.Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        # transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "        # transforms.Normalize([0.2138, 0.0986, 0.0467], [0.0856, 0.0697, 0.0405])\n",
    "        transforms.Normalize([0.3227, 0.2076, 0.1405], [0.0816, 0.0715, 0.0502])\n",
    "    ])\n",
    "\n",
    "col_transform = transforms.Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        # transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "        # transforms.Normalize([0.2138, 0.0986, 0.0467], [0.0856, 0.0697, 0.0405])\n",
    "        transforms.Normalize([0.20489193, 0.07999779, 0.03514091], [0.07887232, 0.06579461, 0.03749502])\n",
    "    ])\n",
    "\n",
    "add_dir = 'img/BED_testdata'\n",
    "BED_add = datasets.ImageFolder(add_dir, data_transform)\n",
    "BED_dl = DataLoader(BED_add, batch_size=1, num_workers=1)\n",
    "print(len(BED_dl))\n",
    "\n",
    "# Taken from the torchvision tutorial\n",
    "# https://pytorch.org/vision/stable/auto_examples/plot_visualization_utils.html\n",
    "model = models.resnext50_32x4d()\n",
    "num_ftrs = model.fc.in_features\n",
    "class_names = ['A', 'B', 'C', 'N']\n",
    "model.fc = torch.nn.Linear(num_ftrs, len(class_names))\n",
    "model.load_state_dict(torch.load('models/final_GT_resnext50_32x4d_finetuned.pth'))\n",
    "model.eval()\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "83ff0cb6-377f-4b04-bb6d-064b4532d39d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted: \"A\", Actual: \"A\", Confidence: 0.93, 0.03, 0.01, 0.03\n",
      "Predicted: \"A\", Actual: \"A\", Confidence: 0.97, 0.01, 0.00, 0.01\n",
      "Predicted: \"A\", Actual: \"A\", Confidence: 0.70, 0.04, 0.01, 0.24\n",
      "Predicted: \"A\", Actual: \"A\", Confidence: 0.66, 0.07, 0.05, 0.22\n",
      "Predicted: \"A\", Actual: \"A\", Confidence: 1.00, 0.00, 0.00, 0.00\n",
      "Predicted: \"B\", Actual: \"B\", Confidence: 0.00, 1.00, 0.00, 0.00\n",
      "Predicted: \"B\", Actual: \"B\", Confidence: 0.00, 1.00, 0.00, 0.00\n",
      "Predicted: \"B\", Actual: \"B\", Confidence: 0.03, 0.80, 0.02, 0.16\n",
      "Predicted: \"B\", Actual: \"B\", Confidence: 0.00, 1.00, 0.00, 0.00\n",
      "Predicted: \"B\", Actual: \"B\", Confidence: 0.04, 0.88, 0.02, 0.06\n",
      "Predicted: \"C\", Actual: \"C\", Confidence: 0.00, 0.00, 0.99, 0.00\n"
     ]
    }
   ],
   "source": [
    "def prettify_prediction(pred):\n",
    "    sm = torch.softmax(pred, dim=1)\n",
    "    lis = [f'{it.item():.2f}' for it in sm[0]]\n",
    "    output = ', '.join(lis)\n",
    "    return output\n",
    "\n",
    "for i, (inputs, labels) in enumerate(BED_dl):\n",
    "    with torch.no_grad():\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "        pred = model(inputs)\n",
    "        predicted, actual = class_names[pred[0].argmax(0)], class_names[labels]\n",
    "        print(f'Predicted: \"{predicted}\", Actual: \"{actual}\", Confidence: {prettify_prediction(pred)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "681eefd1-b596-4e0c-80b6-ea77ba9f0ca1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "target_layers = [model.layer4[-1]]\n",
    "\n",
    "def make_img(inp):\n",
    "    inp = inp.numpy().transpose((1, 2, 0))\n",
    "    mean = np.array([0.3227, 0.2076, 0.1405])\n",
    "    std = np.array([0.0816, 0.0715, 0.0502])\n",
    "    inp = std * inp + mean\n",
    "    inp = np.clip(inp, 0, 1)\n",
    "    return inp\n",
    "\n",
    "def imshow(inp, title=None):\n",
    "    \"\"\"Display image for Tensor.\"\"\"\n",
    "    inp = inp.numpy().transpose((1, 2, 0))\n",
    "    # mean = np.array([0.485, 0.456, 0.406])\n",
    "    # mean = np.array([0.2138, 0.0986, 0.0467])\n",
    "    mean = np.array([0.3227, 0.2076, 0.1405])\n",
    "    # std = np.array([0.229, 0.224, 0.225])\n",
    "    # std = np.array([0.0856, 0.0697, 0.0405])\n",
    "    std = np.array([0.0816, 0.0715, 0.0502])\n",
    "    inp = std * inp + mean\n",
    "    inp = np.clip(inp, 0, 1)\n",
    "    plt.imshow(inp)\n",
    "    if title is not None:\n",
    "        plt.title(title)\n",
    "\n",
    "# GradCAM, HiResCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM, EigenGradCAM, FullGrad\n",
    "\n",
    "sup_classes = ['A', 'B', 'C', 'N']\n",
    "col_classes = ['blue', 'red', 'yellow']\n",
    "\n",
    "def visualize_output(model, class_names, target_layers, data, num_images=8):\n",
    "    cam = HiResCAM(model=model, target_layers=target_layers, use_cuda=True)\n",
    "    images_so_far = 0\n",
    "    cols = 8\n",
    "    rows = int(np.ceil(num_images/cols))\n",
    "    fig = plt.figure(figsize=(int(cols*2.5),int(rows*2.75)))\n",
    "    for i, (inputs, labels) in enumerate(data):\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "        grayscale_cam = cam(input_tensor=inputs)[0]\n",
    "        rgb_img = make_img(inputs[0].cpu())\n",
    "        visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)\n",
    "        outputs = model(inputs)\n",
    "        _, preds = torch.max(outputs, 1)\n",
    "\n",
    "        for j in range(inputs.size()[0]):\n",
    "            images_so_far += 1\n",
    "            ax = plt.subplot(rows, cols, images_so_far)\n",
    "            ax.axis('off')\n",
    "            ax.set_title(f'Predicted: {class_names[preds[j]]}\\n{prettify_prediction(outputs)}')\n",
    "            plt.imshow(visualization)\n",
    "\n",
    "            if images_so_far == num_images:\n",
    "                return\n",
    "\n",
    "def visualize_color_output(model, class_names, target_layers, data, num_images=8):\n",
    "    cam = HiResCAM(model=model, target_layers=target_layers, use_cuda=True)\n",
    "    images_so_far = 0\n",
    "    cols = 4\n",
    "    rows = int(np.ceil(num_images/cols))\n",
    "    fig = plt.figure(figsize=(int(cols*5),int(rows*2.75)))\n",
    "    for i, (inputs, labels) in enumerate(data):\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "        grayscale_cam = cam(input_tensor=inputs)[0]\n",
    "        rgb_img = make_img(inputs[0].cpu())\n",
    "        visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)\n",
    "        normalize_vis = visualization / 256\n",
    "        stacked_imgs = np.hstack([rgb_img, normalize_vis])\n",
    "        outputs = model(inputs)\n",
    "        _, preds = torch.max(outputs, 1)\n",
    "\n",
    "        for j in range(inputs.size()[0]):\n",
    "            images_so_far += 1\n",
    "            ax = plt.subplot(rows, cols, images_so_far)\n",
    "            ax.axis('off')\n",
    "            ax.set_title(f'Predicted: {class_names[preds[j]]}\\n{prettify_prediction(outputs)}')\n",
    "            plt.imshow(stacked_imgs)\n",
    "\n",
    "            if images_so_far == num_images:\n",
    "                return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "54a36136-8a96-4070-87c0-5f2727eeb6e6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 644588 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_output(model, class_names=sup_classes, target_layers=target_layers, num_images=len(BED_dl), data=BED_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bda9ce36-8e04-4a5d-ab39-5311757504df",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'A': 0, 'B': 1, 'C': 2, 'N': 3}\n",
      "884\n"
     ]
    }
   ],
   "source": [
    "BED_train = datasets.ImageFolder('img/GT_square', data_transform)\n",
    "BED_train_dl = DataLoader(BED_train, batch_size=1)\n",
    "\n",
    "print(BED_train.class_to_idx)\n",
    "print(len(list(filter(lambda x: x[1] == 3, BED_train.imgs))))\n",
    "BED_train_A = torch.utils.data.Subset(BED_train, range(0, 200))\n",
    "BED_train_B = torch.utils.data.Subset(BED_train, range(200, 434))\n",
    "BED_train_C = torch.utils.data.Subset(BED_train, range(434, 634))\n",
    "BED_train_N = torch.utils.data.Subset(BED_train, range(634, 1518))\n",
    "\n",
    "BED_dl_A = torch.utils.data.DataLoader(BED_train_A)\n",
    "BED_dl_B = torch.utils.data.DataLoader(BED_train_B)\n",
    "BED_dl_C = torch.utils.data.DataLoader(BED_train_C)\n",
    "BED_dl_N = torch.utils.data.DataLoader(BED_train_N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ce605c6f-9520-4dea-9fda-a9b4f5073107",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1518\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1518"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(len(BED_train))\n",
    "200+234+200+884"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "0f73521a-7d87-480d-9f10-023c8285bbff",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.utils.data.dataset.Subset"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(BED_dl_A.dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "79ba8441-08cb-4311-9d5d-0ccff9efa39e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 323516 characters]",
      "text/plain": [
       "<Figure size 2000x200 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_output(model, class_names=sup_classes, target_layers=target_layers, num_images=8, data=BED_dl_A)\n",
    "plt.savefig('img/documentation/gradcam_A.jpg', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "38daabc9-a84e-4faf-ac9f-4083d7c54741",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 316648 characters]",
      "text/plain": [
       "<Figure size 2000x200 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_output(model, class_names=sup_classes, target_layers=target_layers, num_images=8, data=BED_dl_B)\n",
    "plt.savefig('img/documentation/gradcam_B.jpg', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "6c350bdf-8850-42b9-9ed5-31b20c3da8da",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 292432 characters]",
      "text/plain": [
       "<Figure size 2000x200 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_output(model, class_names=sup_classes, target_layers=target_layers, num_images=8, data=BED_dl_C)\n",
    "plt.savefig('img/documentation/gradcam_C.jpg', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "1683a5d7-9592-420c-8507-08e4da8bc67c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 436912 characters]",
      "text/plain": [
       "<Figure size 2000x200 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_output(model, class_names=sup_classes, target_layers=target_layers, num_images=8, data=BED_dl_N)\n",
    "plt.savefig('img/documentation/gradcam_N.jpg', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6adf41b7-7b3d-4ff9-b92e-4f94dd91251e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "102831\n"
     ]
    }
   ],
   "source": [
    "BE_sig = datasets.ImageFolder('img/BE_signatures/', data_transform)\n",
    "BE_dl = torch.utils.data.DataLoader(BE_sig, shuffle=True)\n",
    "print(len(BE_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "290d700d-a932-4b1c-be84-246bc8d6ebdb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 3051368 characters]",
      "text/plain": [
       "<Figure size 2000x1600 with 48 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_output(model, class_names=sup_classes, target_layers=target_layers, num_images=48, data=BE_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "777049ef-0883-4ac4-ac31-52829c178568",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading GT_colors dataset with 1013 images\n",
      "==> Computing mean and std..\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "273cd87601284018a856310c2e124c30",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1013 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean: [0.20489193 0.07999779 0.03514091] Std: [0.07887232 0.06579461 0.03749502]\n",
      "time elapsed:  25.828746557235718\n",
      "['blue', 'red', 'yellow']\n"
     ]
    }
   ],
   "source": [
    "N_CHANNELS = 3\n",
    "data_dir = 'img/GT_colors/'\n",
    "dataset = datasets.ImageFolder(data_dir, transform=transforms.ToTensor())\n",
    "full_loader = DataLoader(dataset, shuffle=False, num_workers=4)\n",
    "print('Loading GT_colors dataset with', len(dataset), 'images')\n",
    "\n",
    "before = time()\n",
    "mean = torch.zeros(N_CHANNELS)\n",
    "std = torch.zeros(N_CHANNELS)\n",
    "print('==> Computing mean and std..')\n",
    "for inputs, _labels in tqdm(full_loader):\n",
    "    for i in range(N_CHANNELS):\n",
    "        mean[i] += inputs[:,i,:,:].mean()\n",
    "        std[i] += inputs[:,i,:,:].std()\n",
    "mean.div_(len(dataset))\n",
    "std.div_(len(dataset))\n",
    "m_np = mean.numpy()\n",
    "std_np = std.numpy()\n",
    "print('Mean:', m_np, 'Std:', std_np)\n",
    "print(\"time elapsed: \", time()-before)\n",
    "\n",
    "def imshow(inp, title=None):\n",
    "    \"\"\"Display image for Tensor.\"\"\"\n",
    "    inp = inp.numpy().transpose((1, 2, 0))\n",
    "    mean = m_np\n",
    "    std = std_np\n",
    "    inp = std * inp + mean\n",
    "    inp = np.clip(inp, 0, 1)\n",
    "    plt.imshow(inp)\n",
    "    if title is not None:\n",
    "        plt.title(title)\n",
    "\n",
    "data_transform = transforms.Compose([\n",
    "    transforms.Resize(256),\n",
    "    transforms.CenterCrop(224),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(m_np, std_np)\n",
    "])\n",
    "\n",
    "dataset.transform = data_transform\n",
    "class_names = dataset.classes\n",
    "print(class_names)\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e974cfbb-6dab-4976-8dd0-1a671a5077d1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "color_model = models.resnext50_32x4d(weights=None)\n",
    "num_ftrs = color_model.fc.in_features\n",
    "color_model.fc = nn.Linear(num_ftrs, len(col_classes))\n",
    "color_model = color_model.to(device)\n",
    "color_model.load_state_dict(torch.load('models/color-cv/resnext-fold-2.pth'))\n",
    "color_model.eval()\n",
    "\n",
    "color_model.to(device)\n",
    "color_target_layers = [color_model.layer4[-1]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "6358176c-e94a-4d8e-b9c8-ae424203b9d1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6161\n",
      "['A', 'B', 'C']\n"
     ]
    }
   ],
   "source": [
    "BE_supralibros_dir = 'img/BE_supralibros/'\n",
    "BE_sup_ds = datasets.ImageFolder(BE_supralibros_dir, data_transform)\n",
    "BE_sup_dl = torch.utils.data.DataLoader(BE_sup_ds, shuffle=True, num_workers=4)\n",
    "print(len(BE_sup_dl))\n",
    "print(BE_sup_dl.dataset.classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "45688749-9b35-4e95-b45c-70e8777cc276",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11\n",
      "6161\n"
     ]
    }
   ],
   "source": [
    "col_transform = transforms.Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        # transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "        # transforms.Normalize([0.2138, 0.0986, 0.0467], [0.0856, 0.0697, 0.0405])\n",
    "        transforms.Normalize([0.20489193, 0.07999779, 0.03514091], [0.07887232, 0.06579461, 0.03749502])\n",
    "    ])\n",
    "\n",
    "add_dir = 'img/BED_testdata'\n",
    "BED_col_add = datasets.ImageFolder(add_dir, col_transform)\n",
    "BED_col_dl = DataLoader(BED_col_add, batch_size=1, num_workers=1)\n",
    "print(len(BED_col_dl))\n",
    "\n",
    "BE_supralibros_dir = 'img/BE_supralibros/'\n",
    "BE_col_sup_ds = datasets.ImageFolder(BE_supralibros_dir, col_transform)\n",
    "BE_col_sup_dl = torch.utils.data.DataLoader(BE_col_sup_ds, shuffle=True, num_workers=1)\n",
    "print(len(BE_col_sup_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ad718e0c-e9b9-4829-a2c0-bedabc471a53",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted: \"yellow\", Confidence: 0.03, 0.06, 0.91\n",
      "Predicted: \"yellow\", Confidence: 0.00, 0.01, 0.98\n",
      "Predicted: \"yellow\", Confidence: 0.01, 0.03, 0.96\n",
      "Predicted: \"yellow\", Confidence: 0.02, 0.01, 0.98\n",
      "Predicted: \"blue\", Confidence: 1.00, 0.00, 0.00\n",
      "Predicted: \"red\", Confidence: 0.00, 1.00, 0.00\n",
      "Predicted: \"red\", Confidence: 0.00, 0.99, 0.01\n",
      "Predicted: \"red\", Confidence: 0.00, 0.97, 0.03\n",
      "Predicted: \"red\", Confidence: 0.00, 1.00, 0.00\n",
      "Predicted: \"red\", Confidence: 0.06, 0.55, 0.39\n",
      "Predicted: \"red\", Confidence: 0.03, 0.89, 0.08\n"
     ]
    }
   ],
   "source": [
    "for i, (inputs, labels) in enumerate(BED_col_dl):\n",
    "    with torch.no_grad():\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "        pred = color_model(inputs)\n",
    "        predicted = col_classes[pred[0].argmax(0)]\n",
    "        print(f'Predicted: \"{predicted}\", Confidence: {prettify_prediction(pred)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a94ec4f2-ec17-44c5-8c48-d2c594043a6b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 682260 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_output(color_model, class_names=col_classes, target_layers=color_target_layers, num_images=11, data=BED_col_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "13551ddf-7684-4c3d-808d-7dc83844b13a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 900424 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_color_output(color_model, class_names=col_classes, target_layers=color_target_layers, num_images=8, data=BE_col_sup_dl)\n",
    "plt.savefig('img/documentation/gradcam_color.jpg', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23d53c85-347f-4e22-bc9c-857603fdbe81",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
