{
 "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.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": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f03a4bcdad0>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 161908 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "for x, y in loaders['train']:\n",
    "    break\n",
    "    \n",
    "f, arr =  plt.subplots(1, 2, figsize=(10, 5))\n",
    "arr[0].imshow(x[0].permute(1,2,0).detach().cpu().numpy())\n",
    "arr[1].imshow(y[0].detach().cpu().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 10,
     "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_camvid/run1/checkpoint-860.pt\")\n",
    "model.load_state_dict(ckpt[\"state_dict\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f34382ac590>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 144160 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# for x, y in loaders['train']:\n",
    "#     break\n",
    "for x, y in loaders['fine_tune']:\n",
    "    break\n",
    "\n",
    "    \n",
    "with torch.no_grad():\n",
    "    pred = model(x.cuda())\n",
    "    pred = torch.max(pred, dim=1)[1]\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(y[0].detach().cpu().numpy())\n",
    "arr[2].imshow(pred[0].detach().cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 158760 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 196568 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 166400 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# for x, y in loaders['train']:\n",
    "#     break\n",
    "for x, y in loaders['test']:\n",
    "    break\n",
    "\n",
    "    \n",
    "with torch.no_grad():\n",
    "    pred = model(x.cuda())\n",
    "    pred = torch.max(pred, dim=1)[1]\n",
    "    \n",
    "f, arr =  plt.subplots(1, 3, figsize=(15, 5))\n",
    "arr[0].set_title(\"Image\", fontsize=14)\n",
    "arr[0].imshow(x[0].permute(1,2,0).detach().cpu().numpy())\n",
    "arr[1].imshow(y[0].detach().cpu().numpy())\n",
    "arr[1].set_title(\"GT\", fontsize=14)\n",
    "arr[2].imshow(pred[0].detach().cpu().numpy())\n",
    "arr[2].set_title(\"Prediction\", fontsize=14)\n",
    "plt.show()\n",
    "\n",
    "# Transpose\n",
    "\n",
    "x = x.permute(0, 1, 3, 2)\n",
    "y = y.permute(0, 2, 1)\n",
    "\n",
    "with torch.no_grad():\n",
    "    pred = model(x.cuda())\n",
    "    pred = torch.max(pred, dim=1)[1]\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(y[0].detach().cpu().numpy())\n",
    "arr[2].imshow(pred[0].detach().cpu().numpy())\n",
    "plt.show()\n",
    "\n",
    "# Vertical Flip\n",
    "\n",
    "x = x.permute(0, 1, 3, 2)\n",
    "y = y.permute(0, 2, 1)\n",
    "x = x.flip(dims=(-2,))\n",
    "y = y.flip(dims=(-2,))\n",
    "\n",
    "with torch.no_grad():\n",
    "    pred = model(x.cuda())\n",
    "    pred = torch.max(pred, dim=1)[1]\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(y[0].detach().cpu().numpy())\n",
    "arr[2].imshow(pred[0].detach().cpu().numpy())\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Augerino Rotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_train = len(loaders[\"train\"].dataset)\n",
    "n_test = len(loaders[\"test\"].dataset)\n",
    "n_val = len(loaders[\"val\"].dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_rotations(num):\n",
    "    z = torch.randn(num)\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",
    "    return affineMatrices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_rotations = get_rotations(n_train)\n",
    "torch.save(train_rotations, \"train_rotations.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_rotations = get_rotations(n_test)\n",
    "torch.save(test_rotations, \"test_rotations.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_rotations = get_rotations(n_val)\n",
    "torch.save(val_rotations, \"val_rotations.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.savez(\"rotation_tensors.npz\", \n",
    "        train=train_rotations.numpy(),\n",
    "        test=test_rotations.numpy(),\n",
    "        val=val_rotations.numpy(),\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading a dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from augerino.rot_camvid_data import RotCamVid, rot_camvid_loaders\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.load(\"rotation_tensors.npz\")\n",
    "train_rotations = torch.from_numpy(arr[\"train\"]).float()\n",
    "test_rotations = torch.from_numpy(arr[\"test\"]).float()\n",
    "val_rotations = torch.from_numpy(arr[\"val\"]).float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/izmailovpavel/datasets/CamVid/\n"
     ]
    }
   ],
   "source": [
    "loaders, num_classes = rot_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",
    "    train_rotations=train_rotations,\n",
    "    test_rotations=test_rotations,\n",
    "    val_rotations=val_rotations,\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": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f22b0248590>"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 142568 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "for x, y in loaders['train']:\n",
    "    break\n",
    "    \n",
    "f, arr =  plt.subplots(1, 2, figsize=(10, 5))\n",
    "arr[0].imshow(x[0].permute(1,2,0).detach().cpu().numpy())\n",
    "arr[1].imshow(y[0].detach().cpu().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### For a given image rotation is consistent\n",
    "\n",
    "Train loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f22842ba350>"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 151220 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x, y = loaders['train'].dataset[100]\n",
    "    \n",
    "f, arr =  plt.subplots(1, 2, figsize=(10, 5))\n",
    "arr[0].imshow(x.permute(1,2,0).detach().cpu().numpy())\n",
    "arr[1].imshow(y.detach().cpu().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fine-Tune loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f2275cce450>"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 150224 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x, y = loaders['fine_tune'].dataset[100]\n",
    "    \n",
    "f, arr =  plt.subplots(1, 2, figsize=(10, 5))\n",
    "arr[0].imshow(x.permute(1,2,0).detach().cpu().numpy())\n",
    "arr[1].imshow(y.detach().cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f2275e64890>"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 144440 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x, y = loaders['test'].dataset[10]\n",
    "    \n",
    "f, arr =  plt.subplots(1, 2, figsize=(10, 5))\n",
    "arr[0].imshow(x.permute(1,2,0).detach().cpu().numpy())\n",
    "arr[1].imshow(y.detach().cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f2275b78dd0>"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 151892 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x, y = loaders['val'].dataset[10]\n",
    "    \n",
    "f, arr =  plt.subplots(1, 2, figsize=(10, 5))\n",
    "arr[0].imshow(x.permute(1,2,0).detach().cpu().numpy())\n",
    "arr[1].imshow(y.detach().cpu().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Affine Augmentations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from augerino.aug_eq_model import UniformEqAug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "aug = UniformEqAug(padding=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f0410bd3690>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 153688 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "for x, y in loaders['test']:\n",
    "    break\n",
    "    \n",
    "f, arr =  plt.subplots(1, 2, figsize=(10, 5))\n",
    "arr[0].imshow(x[0].permute(1,2,0).detach().cpu().numpy())\n",
    "arr[1].imshow(y[0].detach().cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).\n",
      "Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f041081de10>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 115368 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "f, arr =  plt.subplots(1, 4, figsize=(10, 5))\n",
    "\n",
    "arr[0].imshow(x[0].permute(1,2,0).detach().cpu().numpy())\n",
    "arr[1].imshow(aug(x)[0].permute(1,2,0).detach().cpu().numpy())\n",
    "arr[2].imshow(aug.inverse(aug(x))[0].permute(1,2,0).detach().cpu().numpy())\n",
    "arr[3].imshow(aug.inverse(x)[0].permute(1,2,0).detach().cpu().numpy())"
   ]
  }
 ],
 "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
}
