{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "sys.path.insert(0, \"..\")\n",
    "from deployment.handler import TwinHandler, HookCAM, HookCAMBwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Recreate Model and Load Weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "twin_handler = TwinHandler()\n",
    "image_tfm = twin_handler.image_tfm\n",
    "\n",
    "head_reload = twin_handler.head_reload\n",
    "encoder_reload, _ = twin_handler.get_encoder(pre_train=False)\n",
    "type(encoder_reload), type(head_reload)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_head = torch.load(\"../model/resnet50_0.962_head.pth\")\n",
    "head_reload.load_state_dict(state_head)\n",
    "head_reload.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_encoder = torch.load(\"../model/resnet50_0.962_encoder.pth\")\n",
    "encoder_reload.load_state_dict(state_encoder)\n",
    "encoder_reload.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch Inference and CAM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imgtest, imgval = (\n",
    "    Image.open(\"../sample/c3.jpg\").convert(\"RGB\"),\n",
    "    Image.open(\"../sample/c4.jpg\").convert(\"RGB\"),\n",
    ")\n",
    "\n",
    "_, axes = plt.subplots(1, 2, figsize=(10, 5))\n",
    "axes[0].imshow(imgtest)\n",
    "axes[0].axis(\"off\")\n",
    "axes[1].imshow(imgval)\n",
    "axes[1].axis(\"off\")\n",
    "plt.show()\n",
    "\n",
    "imgtest_ts, imgval_ts = image_tfm(imgtest)[None, ...], image_tfm(imgval)[None, ...]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def img_restore(image: torch.Tensor):\n",
    "        return (image - image.min()) / (image.max() - image.min())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "with HookCAMBwd(encoder_reload) as hookg:\n",
    "    with HookCAM(encoder_reload) as hook:\n",
    "        \n",
    "        encoder_reload.to(torch.device(\"cuda\"))\n",
    "        l_emb = encoder_reload(imgtest_ts.cuda())\n",
    "        r_emb = encoder_reload(imgval_ts.cuda())\n",
    "\n",
    "        ftrs = torch.cat([l_emb, r_emb], dim=1)\n",
    "        head_reload.to(torch.device(\"cuda\"))\n",
    "        res = head_reload(ftrs)[0]\n",
    "        act = hook.stored\n",
    "\n",
    "    pred_cls = res.argmax().item()\n",
    "    res[pred_cls].backward()\n",
    "    grad = hookg.stored\n",
    "    \n",
    "encoder_reload.zero_grad(), head_reload.zero_grad()\n",
    "pred_cls, [\"Not Similar\", \"Similar\"][pred_cls]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weight_left = grad[0][0].mean(dim=[1, 2], keepdim=True)\n",
    "cam_map_left = (weight_left * act[0][0]).sum(0)\n",
    "\n",
    "weight_right = grad[1][0].mean(dim=[1, 2], keepdim=True)\n",
    "cam_map_right = (weight_right * act[1][0]).sum(0)\n",
    "print(cam_map_left.shape, cam_map_right.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, ax = plt.subplots(1, 2, figsize=(10, 5))\n",
    "ax[0].imshow(img_restore(imgtest_ts[0].permute(1, 2, 0)))\n",
    "ax[0].imshow(\n",
    "    cam_map_left.detach().cpu(),\n",
    "    alpha=0.3,\n",
    "    extent=(0, 224, 224, 0),\n",
    "    interpolation=\"bilinear\",\n",
    "    cmap=\"jet\",\n",
    ")\n",
    "ax[0].axis(\"off\")\n",
    "\n",
    "ax[1].imshow(img_restore(imgval_ts[0].permute(1, 2, 0)))\n",
    "ax[1].imshow(\n",
    "    cam_map_right.detach().cpu(),\n",
    "    alpha=0.3,\n",
    "    extent=(0, 224, 224, 0),\n",
    "    interpolation=\"bilinear\",\n",
    "    cmap=\"jet\",\n",
    ")\n",
    "ax[1].axis(\"off\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TorchServe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### handle.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pygmentize ../deployment/handler.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Archive Model: Encoder and Head"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```bash\n",
    "cd ../\n",
    "\n",
    "cp model/resnet50_0.962_head.pth model/head_weight.pth\n",
    "cp model/resnet50_0.962_encoder.pth model/encoder_weight.pth\n",
    "\n",
    "torch-model-archiver --model-name twin --version 1.0 --serialized-file ./model/encoder_weight.pth --export-path model_store --handler ./deployment/handler.py -f --extra-files ./model/head_weight.pth\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Serve the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```bash\n",
    "torchserve --start --ncs --model-store model_store --models twin.mar\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Call API"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```bash\n",
    "time http --form POST http://127.0.0.1:8080/predictions/twin left@sample/c1.jpg right@sample/c2.jpg cam=False\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample Response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```bash\n",
    "HTTP/1.1 200\n",
    "Cache-Control: no-cache; no-store, must-revalidate, private\n",
    "Expires: Thu, 01 Jan 1970 00:00:00 UTC\n",
    "Pragma: no-cache\n",
    "connection: keep-alive\n",
    "content-length: 46\n",
    "x-request-id: e51e2f15-a9b6-4522-b700-c0f5e35008a7\n",
    "\n",
    "[\n",
    "  -0.938737690448761,\n",
    "  0.7865392565727234\n",
    "]\n",
    "\n",
    "\n",
    "real    0m1.765s\n",
    "user    0m0.334s\n",
    "sys     0m0.036s\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stop TorchServe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```sh\n",
    "torchserve --stop\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "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.10"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}