{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import os, sys\n",
    "import argparse\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "from functools import partial\n",
    "\n",
    "import utils.training as train_utils\n",
    "\n",
    "from utils.training import adjust_learning_rate, schedule, save_checkpoint\n",
    "from utils.training import seg_cross_entropy\n",
    "import augerino.models as models\n",
    "from augerino.camvid_data import camvid_loaders\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_cfg = models.FCDenseNet67"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/izmailovpavel/datasets/CamVid/\n"
     ]
    }
   ],
   "source": [
    "loaders, num_classes = camvid_loaders(\n",
    "    \"/home/izmailovpavel/datasets/CamVid/\",\n",
    "    4,\n",
    "    4,\n",
    "    ft_batch_size=1,\n",
    "    transform_train=model_cfg.transform_train,\n",
    "    transform_test=model_cfg.transform_test,\n",
    "    joint_transform=model_cfg.joint_transform,\n",
    "    ft_joint_transform=model_cfg.ft_joint_transform,\n",
    "    target_transform=model_cfg.target_transform\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = model_cfg.base(\n",
    "    *model_cfg.args,\n",
    "    num_classes=num_classes,\n",
    "    **model_cfg.kwargs,\n",
    "    use_aleatoric=False\n",
    ")\n",
    "model.cuda();\n",
    "\n",
    "ckpt = torch.load(\"ckpts/tiramisu_rotcamvid/run1/checkpoint-900.pt\")\n",
    "model.load_state_dict(ckpt[\"state_dict\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rotate(img, z):\n",
    "    num = img.shape[0]\n",
    "    affineMatrices = torch.zeros(num, 2, 3)\n",
    "    affineMatrices[:,0,0] = z.cos()\n",
    "    affineMatrices[:,0,1] = -z.sin()\n",
    "    affineMatrices[:,1,0] = z.sin()\n",
    "    affineMatrices[:,1,1] = z.cos()\n",
    "    \n",
    "    if len(img.shape) == 3:\n",
    "        img_ = img.float()[:, None] + 1\n",
    "        flowgrid = F.affine_grid(affineMatrices, size=img_.size(), align_corners=True)\n",
    "        img_ = F.grid_sample(img_, flowgrid, align_corners=True).int()\n",
    "        img_ -= 1\n",
    "        img_[img_ == -1] = 12\n",
    "        return img_[:, 0]\n",
    "    else:\n",
    "        flowgrid = F.affine_grid(affineMatrices, size=img.size(), align_corners=True)\n",
    "        return F.grid_sample(img, flowgrid, align_corners=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "for x, y in loaders[\"test\"]:\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 844420 characters]\n",
      "text/plain": [
       "<Figure size 1080x1800 with 15 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "zs = torch.linspace(0, np.pi, 5)\n",
    "\n",
    "nrow, ncol = len(zs), 3\n",
    "f, arr = plt.subplots(nrow, ncol, figsize=(ncol*5, nrow*5))\n",
    "\n",
    "for i, z in enumerate(zs):\n",
    "    x_ = rotate(x, z[None].repeat((4)))\n",
    "    y_ = rotate(y, z[None].repeat((4)))\n",
    "    with torch.no_grad():\n",
    "        pred = model(x_.cuda())\n",
    "        pred = torch.max(pred, dim=1)[1]\n",
    "    arr[i, 0].imshow(x_[0].permute(1,2,0).detach().cpu().numpy())\n",
    "    arr[i, 1].imshow(y_[0].detach().cpu().numpy())\n",
    "    arr[i, 2].imshow(pred[0].detach().cpu().numpy())\n",
    "    \n",
    "arr[0, 0].imshow(x[0].permute(1,2,0).detach().cpu().numpy())\n",
    "arr[0, 1].set_title(\"GT\", fontsize=14)\n",
    "arr[0, 2].imshow(pred[0].detach().cpu().numpy())\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## stabilized view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 916092 characters]\n",
      "text/plain": [
       "<Figure size 1440x1800 with 20 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "zs = torch.linspace(0, np.pi, 5)\n",
    "\n",
    "nrow, ncol = len(zs), 4\n",
    "f, arr = plt.subplots(nrow, ncol, figsize=(ncol*5, nrow*5))\n",
    "\n",
    "for i, z in enumerate(zs):\n",
    "    x_ = rotate(x, z[None].repeat((4)))\n",
    "    y_ = rotate(y, z[None].repeat((4)))\n",
    "    with torch.no_grad():\n",
    "        pred = model(x_.cuda())\n",
    "        pred = torch.max(pred, dim=1)[1]\n",
    "    pred_ = rotate(pred.detach().cpu(), -z[None].repeat((4)))\n",
    "    arr[i, 0].imshow(x_[0].permute(1,2,0).detach().cpu().numpy())\n",
    "    arr[i, 1].imshow(y_[0].detach().cpu().numpy())\n",
    "    arr[i, 2].imshow(pred[0].detach().cpu().numpy())\n",
    "    arr[i, 3].imshow(pred_[0].detach().cpu().numpy())\n",
    "    \n",
    "arr[0, 0].set_title(\"Image\", fontsize=14)\n",
    "arr[0, 1].set_title(\"GT\", fontsize=14)\n",
    "arr[0, 2].set_title(\"Prediction\", fontsize=14)\n",
    "arr[0, 3].set_title(\"Stabilized\", fontsize=14)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's make a rotationally-equivariant version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from augerino.aug_eq_model import DiffEqAug, AugEqModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "aug = DiffEqAug(padding=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f72ffdb9d10>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 362844 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "for x, y in loaders['test']:\n",
    "    break\n",
    "    \n",
    "with torch.no_grad():\n",
    "    x1 = aug(x)\n",
    "    x2 = aug.inverse(x1)\n",
    "    \n",
    "f, arr =  plt.subplots(1, 3, figsize=(15, 5))\n",
    "arr[0].imshow(x[0].permute(1,2,0).detach().cpu().numpy())\n",
    "arr[1].imshow(x1[0].permute(1,2,0).detach().cpu().numpy())\n",
    "arr[2].imshow(x2[0].permute(1,2,0).detach().cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "aug = DiffEqAug(padding=20)\n",
    "augerino_model = AugEqModel(model, aug)\n",
    "augerino_model.cuda();\n",
    "augerino_model.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 899820 characters]\n",
      "text/plain": [
       "<Figure size 1440x1800 with 20 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "zs = torch.linspace(0, np.pi, 5)\n",
    "\n",
    "nrow, ncol = len(zs), 4\n",
    "f, arr = plt.subplots(nrow, ncol, figsize=(ncol*5, nrow*5))\n",
    "\n",
    "for i, z in enumerate(zs):\n",
    "    x_ = rotate(x, z[None].repeat((4)))\n",
    "    y_ = rotate(y, z[None].repeat((4)))\n",
    "    with torch.no_grad():\n",
    "        pred = augerino_model(x_.cuda())\n",
    "        pred = torch.max(pred, dim=1)[1]\n",
    "    pred_ = rotate(pred.detach().cpu(), -z[None].repeat((4)))\n",
    "    arr[i, 0].imshow(x_[0].permute(1,2,0).detach().cpu().numpy())\n",
    "    arr[i, 1].imshow(y_[0].detach().cpu().numpy())\n",
    "    arr[i, 2].imshow(pred[0].detach().cpu().numpy())\n",
    "    arr[i, 3].imshow(pred_[0].detach().cpu().numpy())\n",
    "    \n",
    "arr[0, 0].set_title(\"Image\", fontsize=14)\n",
    "arr[0, 1].set_title(\"GT\", fontsize=14)\n",
    "arr[0, 2].set_title(\"Prediction\", fontsize=14)\n",
    "arr[0, 3].set_title(\"Stabilized\", fontsize=14)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Augerino Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "aug = DiffEqAug(padding=20)\n",
    "augerino_model = AugEqModel(model, aug)\n",
    "augerino_model.cuda();\n",
    "augerino_model.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ckpt = torch.load(\"ckpts/tiramisu_rotcamvid/augerino_run1/checkpoint-300.pt\")\n",
    "augerino_model.load_state_dict(ckpt[\"state_dict\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 898660 characters]\n",
      "text/plain": [
       "<Figure size 1440x1800 with 20 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "zs = torch.linspace(0, np.pi, 5)\n",
    "\n",
    "nrow, ncol = len(zs), 4\n",
    "f, arr = plt.subplots(nrow, ncol, figsize=(ncol*5, nrow*5))\n",
    "\n",
    "for i, z in enumerate(zs):\n",
    "    x_ = rotate(x, z[None].repeat((4)))\n",
    "    y_ = rotate(y, z[None].repeat((4)))\n",
    "    with torch.no_grad():\n",
    "        pred = augerino_model(x_.cuda())\n",
    "        pred = torch.max(pred, dim=1)[1]\n",
    "    pred_ = rotate(pred.detach().cpu(), -z[None].repeat((4)))\n",
    "    arr[i, 0].imshow(x_[0].permute(1,2,0).detach().cpu().numpy())\n",
    "    arr[i, 1].imshow(y_[0].detach().cpu().numpy())\n",
    "    arr[i, 2].imshow(pred[0].detach().cpu().numpy())\n",
    "    arr[i, 3].imshow(pred_[0].detach().cpu().numpy())\n",
    "    \n",
    "arr[0, 0].set_title(\"Image\", fontsize=14)\n",
    "arr[0, 1].set_title(\"GT\", fontsize=14)\n",
    "arr[0, 2].set_title(\"Prediction\", fontsize=14)\n",
    "arr[0, 3].set_title(\"Stabilized\", fontsize=14)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py37",
   "language": "python",
   "name": "py37"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
