modules.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:51.437487Z",
     "start_time": "2020-08-06T18:57:51.435893Z"
    }
   },
   "outputs": [],
   "source": [
    "# default_exp modules"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This contains pytorch modules which allow for optimization of intrinsic/extrinsic parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-14T13:42:09.485783Z",
     "start_time": "2020-05-14T13:42:09.141697Z"
    }
   },
   "source": [
    "# Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.448530Z",
     "start_time": "2020-08-06T18:57:51.443396Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from camera_calib.utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.452118Z",
     "start_time": "2020-08-06T18:57:52.449815Z"
    }
   },
   "outputs": [],
   "source": [
    "from IPython.core.debugger import set_trace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Utility"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.466062Z",
     "start_time": "2020-08-06T18:57:52.453366Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def tensors2parameters(args): return args_loop(args, lambda arg: nn.Parameter(arg))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-30T14:20:52.666739Z",
     "start_time": "2020-04-30T14:20:52.665136Z"
    }
   },
   "source": [
    "# Inversible Modules"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some things:\n",
    "* I want to be able to have a `Module`, call `Inverse()` on it, and have the `Module` and it's `Inverse` share the same parameters, yet the inverse does the inverse of the module\n",
    "* I want to ensure some intuitive things, like: `Inverse(Inverse(Module)) == Module`\n",
    "* It's useful to be able to access some parameters in an \"inverse-aware\" reconstructed form, so I've implemented `get_param()`. i.e. if a `Rotation` is `Inverse`d and I want to get the rotation matrix, calling the `get_param()` method should return the transposed version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For inversible modules, inherit this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.472352Z",
     "start_time": "2020-08-06T18:57:52.467181Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Inversible(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "       \n",
    "    def forward_param(self): raise NotImplementedError('Please implement forward_param() method')\n",
    "    def inverse_param(self): raise NotImplementedError('Please implement inverse_param() method')\n",
    "        \n",
    "    # Call this method to get a copy of \"inverse aware\" reconstructed parameters\n",
    "    def get_param(self):\n",
    "        with torch.no_grad():\n",
    "            return self.forward_param().detach().clone() \n",
    "    \n",
    "    def inverse(self): raise NotImplementedError('Please implement inverse() method')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-30T03:51:39.031196Z",
     "start_time": "2020-04-30T03:51:39.027408Z"
    }
   },
   "source": [
    "Do `Inverse(Module)` to take its inverse:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.478628Z",
     "start_time": "2020-08-06T18:57:52.473488Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Inverse(Inversible): # Note that Inverse is itself Inversible...\n",
    "    def __init__(self, m):\n",
    "        super().__init__()\n",
    "        self.m = m\n",
    "        \n",
    "    def __repr__(self):\n",
    "        return f'{self.__class__.__name__}({self.m.__repr__()})'\n",
    "    \n",
    "    def forward_param(self): return self.m.inverse_param()\n",
    "    def inverse_param(self): return self.m.forward_param()\n",
    "    \n",
    "    def forward(self, x): return self.m.inverse(x)\n",
    "    def inverse(self, x): return self.m.forward(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Have a basic test to test inversibility"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.487224Z",
     "start_time": "2020-08-06T18:57:52.479612Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def assert_inversible(m, x, y, **kwargs):\n",
    "    assert_allclose(m(x),          y, **kwargs)\n",
    "    assert_allclose(Inverse(m)(y), x, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.494953Z",
     "start_time": "2020-08-06T18:57:52.488148Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Inversible(), Inverse(Inversible()), Inverse(Inverse(Inversible())))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = Inversible(); b = Inverse(a); c = Inverse(b)\n",
    "a,b,c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Translation module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-02T01:21:28.774565Z",
     "start_time": "2020-05-02T01:21:28.768732Z"
    }
   },
   "source": [
    "This performs a simple 3D translation; the inverse is simply subtracting the translation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.501046Z",
     "start_time": "2020-08-06T18:57:52.496426Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Translation(Inversible):\n",
    "    def __init__(self, t):\n",
    "        super().__init__()\n",
    "        self.t = tensors2parameters(t)\n",
    "        \n",
    "    def __repr__(self):\n",
    "        t_x, t_y, t_z = self.t\n",
    "        return f'{self.__class__.__name__}(tx:{t_x:.4} ty:{t_y:.4} tz:{t_z:.4})'\n",
    "        \n",
    "    def forward_param(self): return  self.t\n",
    "    def inverse_param(self): return -self.forward_param()\n",
    "        \n",
    "    def forward(self, ps, inverse=False): \n",
    "        t = self.forward_param() if not inverse else self.inverse_param()\n",
    "        return ps + t\n",
    "    def inverse(self, ps): return self.forward(ps, inverse=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.508697Z",
     "start_time": "2020-08-06T18:57:52.502367Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Translation(tx:1.0 ty:2.0 tz:3.0), Inverse(Translation(tx:1.0 ty:2.0 tz:3.0)))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t = torch.FloatTensor([1,2,3])\n",
    "a = Translation(t); b = Inverse(a);\n",
    "a, b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.515857Z",
     "start_time": "2020-08-06T18:57:52.509812Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0754, 0.0018, 0.8717],\n",
       "        [0.3523, 0.7136, 0.2060],\n",
       "        [0.4739, 0.1833, 0.0573],\n",
       "        [0.7991, 0.4707, 0.4526]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ps = torch.rand(4,3)\n",
    "ps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.520888Z",
     "start_time": "2020-08-06T18:57:52.516957Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_inversible(a, ps, ps+t, atol=1e-4)\n",
    "assert_inversible(b, ps, ps-t, atol=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.526854Z",
     "start_time": "2020-08-06T18:57:52.522083Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose(a.get_param(),  t)\n",
    "assert_allclose(b.get_param(), -t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rotation module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This performs a simple 3D rotation; the inverse is simply the transpose. Implement the reconstruction in the `r2R()` method for child classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.532813Z",
     "start_time": "2020-08-06T18:57:52.527943Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Rotation(Inversible):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "    def r2R(self):\n",
    "        raise NotImplementedError('Please implement rotation matrix reconstruction')       \n",
    "    \n",
    "    def forward_param(self): return self.r2R()\n",
    "    def inverse_param(self): return self.forward_param().T\n",
    "        \n",
    "    def forward(self, ps, inverse=False): \n",
    "        R = self.forward_param() if not inverse else self.inverse_param()\n",
    "        return pmm(ps, R)\n",
    "    def inverse(self, ps): return self.forward(ps, inverse=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Euler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.541760Z",
     "start_time": "2020-08-06T18:57:52.533773Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class EulerRotation(Rotation):\n",
    "    def __init__(self, R):\n",
    "        super().__init__()\n",
    "        self.euler = tensors2parameters(R2euler(R))\n",
    "    \n",
    "    def __repr__(self):\n",
    "        e_x, e_y, e_z = self.euler\n",
    "        return f'{self.__class__.__name__}(ex:{e_x:.4} ey:{e_y:.4} ez:{e_z:.4})'\n",
    "    \n",
    "    def r2R(self): return euler2R(self.euler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.549559Z",
     "start_time": "2020-08-06T18:57:52.543025Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(EulerRotation(ex:-2.142 ey:1.142 ez:-0.1416),\n",
       " Inverse(EulerRotation(ex:-2.142 ey:1.142 ez:-0.1416)))"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "R = euler2R(torch.FloatTensor([1,2,3]))\n",
    "a = EulerRotation(R); b = Inverse(a)\n",
    "a,b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.555131Z",
     "start_time": "2020-08-06T18:57:52.550729Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.7089, 0.4242, 0.6398],\n",
       "        [0.6279, 0.2782, 0.2785],\n",
       "        [0.0057, 0.9793, 0.1631],\n",
       "        [0.6987, 0.9158, 0.2444]])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ps = torch.rand(4,3)\n",
    "ps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.562497Z",
     "start_time": "2020-08-06T18:57:52.556410Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_inversible(a, ps, pmm(ps, R  ), atol=1e-4)\n",
    "assert_inversible(b, ps, pmm(ps, R.T), atol=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.566411Z",
     "start_time": "2020-08-06T18:57:52.563586Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose(a.get_param(), R)\n",
    "assert_allclose(b.get_param(), R.T)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# InversibleSequential Module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having inversible composible modules will be super useful. This is basically nn.Sequential but I've made it inversible. Note that:\n",
    "`y = f(g(h(x)))` implies that `x = h^-1(g^-1(f^-1(y)))`. So basically just reverse the order and take the inverse of each function to invert the whole thing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.571496Z",
     "start_time": "2020-08-06T18:57:52.567390Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class InversibleSequential(Inversible):\n",
    "    def __init__(self, ms):\n",
    "        super().__init__()\n",
    "        self.ms = nn.ModuleList(ms)  \n",
    "    \n",
    "    def forward(self, x):\n",
    "        for m in self.ms: \n",
    "            x = m.forward(x)\n",
    "        return x\n",
    "    \n",
    "    def inverse(self, x):\n",
    "        for m in reversed(self.ms):\n",
    "            x = m.inverse(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rigid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Rigid transform is a rotation followed by translation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.577655Z",
     "start_time": "2020-08-06T18:57:52.572768Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Rigid(InversibleSequential):\n",
    "    def __init__(self, R, t, Rotation=EulerRotation):\n",
    "        super().__init__([Rotation(R), Translation(t)]) # NOTE: rotation happens first\n",
    "        \n",
    "    def forward_param(self):\n",
    "        return Rt2M(self.ms[0].forward_param(), # R \n",
    "                    self.ms[1].forward_param()) # t\n",
    "        \n",
    "    def inverse_param(self): return invert_rigid(self.forward_param())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.585135Z",
     "start_time": "2020-08-06T18:57:52.578647Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Rigid(\n",
       "   (ms): ModuleList(\n",
       "     (0): EulerRotation(ex:-2.142 ey:1.142 ez:-0.1416)\n",
       "     (1): Translation(tx:1.0 ty:2.0 tz:3.0)\n",
       "   )\n",
       " ), Inverse(Rigid(\n",
       "   (ms): ModuleList(\n",
       "     (0): EulerRotation(ex:-2.142 ey:1.142 ez:-0.1416)\n",
       "     (1): Translation(tx:1.0 ty:2.0 tz:3.0)\n",
       "   )\n",
       " )))"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "R, t = euler2R(tuple(torch.FloatTensor([1,2,3]))), torch.FloatTensor([1,2,3])\n",
    "a = Rigid(R, t); b = Inverse(a)\n",
    "a,b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.592825Z",
     "start_time": "2020-08-06T18:57:52.586037Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.8939, 0.3311, 0.8731],\n",
       "        [0.5188, 0.2413, 0.7948],\n",
       "        [0.7772, 0.3186, 0.5079],\n",
       "        [0.0183, 0.1378, 0.0013]])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ps = torch.rand(4,3)\n",
    "ps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.597909Z",
     "start_time": "2020-08-06T18:57:52.595468Z"
    }
   },
   "outputs": [],
   "source": [
    "M = Rt2M(R,t)\n",
    "M_inv = invert_rigid(M)\n",
    "R_inv, t_inv = M2Rt(M_inv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.605949Z",
     "start_time": "2020-08-06T18:57:52.599368Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_inversible(a, ps, pmm(ps, R    ) + t,     atol=1e-4)\n",
    "assert_inversible(b, ps, pmm(ps, R_inv) + t_inv, atol=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.611341Z",
     "start_time": "2020-08-06T18:57:52.607908Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose(a.get_param(), M)\n",
    "assert_allclose(b.get_param(), M_inv)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-02T03:08:06.759294Z",
     "start_time": "2020-05-02T03:08:06.752916Z"
    }
   },
   "source": [
    "Create a wrapper for composing multiple `Rigid`s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.615647Z",
     "start_time": "2020-08-06T18:57:52.612256Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Rigids(InversibleSequential):\n",
    "    def __init__(self, rigids):\n",
    "        super().__init__(rigids)\n",
    "    \n",
    "    def forward_param(self):\n",
    "        M = self.ms[0].forward_param()\n",
    "        for m in self.ms[1:]:\n",
    "            M = mult_rigid(m.forward_param(), M)\n",
    "        return M\n",
    "    \n",
    "    def inverse_param(self): return invert_rigid(self.forward_param())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.623256Z",
     "start_time": "2020-08-06T18:57:52.616566Z"
    }
   },
   "outputs": [],
   "source": [
    "R_a, t_a = euler2R(torch.FloatTensor([ 1, 2, 3])), torch.FloatTensor([1,2,3])\n",
    "a = Rigid(R_a, t_a)\n",
    "R_b, t_b = euler2R(torch.FloatTensor([.5, 1, 1])), torch.FloatTensor([3,2,2])\n",
    "b = Inverse(Rigid(R_b, t_b))\n",
    "R_c, t_c = euler2R(torch.FloatTensor([.3,.1, 0])), torch.FloatTensor([3,3,3])\n",
    "c = Rigid(R_c, t_c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.626989Z",
     "start_time": "2020-08-06T18:57:52.624431Z"
    }
   },
   "outputs": [],
   "source": [
    "d = Rigids([a,b,c])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.632958Z",
     "start_time": "2020-08-06T18:57:52.627913Z"
    }
   },
   "outputs": [],
   "source": [
    "M_a = Rt2M(R_a, t_a)\n",
    "M_b = Rt2M(R_b.T, -R_b.T@t_b)\n",
    "M_c = Rt2M(R_c, t_c)\n",
    "M_d = M_c@M_b@M_a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.638130Z",
     "start_time": "2020-08-06T18:57:52.634096Z"
    }
   },
   "outputs": [],
   "source": [
    "R_d, t_d = M2Rt(M_d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.647524Z",
     "start_time": "2020-08-06T18:57:52.639295Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_inversible(d, ps, pmm(ps, R_d) + t_d, atol=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.651957Z",
     "start_time": "2020-08-06T18:57:52.648733Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose(d.get_param(), M_d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-05T21:55:43.455346Z",
     "start_time": "2020-05-05T21:55:43.453510Z"
    }
   },
   "source": [
    "# Normalize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will normalize 3D coordinates (i.e. project points to \"unit\" image plane)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.656502Z",
     "start_time": "2020-08-06T18:57:52.652809Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Normalize(nn.Module): # Note: Not Inversible\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "    def forward(self, ps): return normalize(ps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.664058Z",
     "start_time": "2020-08-06T18:57:52.657399Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.4744, 0.9574, 0.2313],\n",
       "        [0.9880, 0.3148, 0.3756],\n",
       "        [0.3908, 0.1181, 0.0731],\n",
       "        [0.0018, 0.5556, 0.5136]])"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ps = torch.rand(4,3)\n",
    "ps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.668741Z",
     "start_time": "2020-08-06T18:57:52.665205Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose(Normalize()(ps), ps[:,0:2]/ps[:,[2]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Augment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-05T23:26:15.800252Z",
     "start_time": "2020-05-05T23:26:15.791287Z"
    }
   },
   "source": [
    "Adds `one`s to input array; note this is `Inversible`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.675139Z",
     "start_time": "2020-08-06T18:57:52.669767Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Augment(Inversible):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, ps): return augment(ps)\n",
    "    def inverse(self, ps): return deaugment(ps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test to make sure `one`s are added and removed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.685846Z",
     "start_time": "2020-08-06T18:57:52.676620Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.8044, 0.5761],\n",
       "        [0.7281, 0.5730],\n",
       "        [0.6866, 0.4281],\n",
       "        [0.2840, 0.2308]])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ps = torch.rand(4,2)\n",
    "ps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.694341Z",
     "start_time": "2020-08-06T18:57:52.687293Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_inversible(Augment(), ps, torch.cat([ps, ps.new_ones(len(ps),1)], dim=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distortion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will apply distortion model to 2D points. Note for now I didn't find much commonality so I didn't make a base class. Also note that distortions are typically not inversible (at least not analytically to my knowledge), so exclude for now."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### No distortion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.703222Z",
     "start_time": "2020-08-06T18:57:52.695580Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "NoDistortion = nn.Identity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-05T23:35:04.096552Z",
     "start_time": "2020-05-05T23:35:04.092319Z"
    }
   },
   "source": [
    "### Heikkila97 distortion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.712554Z",
     "start_time": "2020-08-06T18:57:52.704305Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Heikkila97Distortion(nn.Module):\n",
    "    def __init__(self, d):\n",
    "        super().__init__()\n",
    "        self.d = tensors2parameters(d)\n",
    "    \n",
    "    def __repr__(self):\n",
    "        k1, k2, p1, p2 = self.d\n",
    "        return f'{self.__class__.__name__}(k1:{k1:.4} k2:{k2:.4} p1:{p1:.4} p2:{p2:.4})'\n",
    "    \n",
    "    def forward(self, ps):\n",
    "        k1, k2, p1, p2 = self.d\n",
    "        xs, ys = ps[:,0], ps[:,1]\n",
    "        \n",
    "        # Radial distortion\n",
    "        rs = xs**2 + ys**2\n",
    "        xs_r = xs*(1 + k1*rs + k2*rs**2)\n",
    "        ys_r = ys*(1 + k1*rs + k2*rs**2)\n",
    "\n",
    "        # Decentering distortion \n",
    "        xs_d = xs_r + 2*p1*xs*ys + p2*(3*xs**2 + ys**2)\n",
    "        ys_d = ys_r + p1*(xs**2 + 3*ys**2) + 2*p2*xs*ys\n",
    "        \n",
    "        return stackify((xs_d, ys_d), dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.722177Z",
     "start_time": "2020-08-06T18:57:52.713589Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Heikkila97Distortion(k1:0.01 k2:0.02 p1:0.03 p2:0.04)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "distortion = Heikkila97Distortion(torch.FloatTensor([.01, 0.02, 0.03, 0.04]))\n",
    "distortion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.729018Z",
     "start_time": "2020-08-06T18:57:52.723539Z"
    }
   },
   "outputs": [],
   "source": [
    "ps = torch.FloatTensor([[0.1900, 0.2467],\n",
    "                        [0.9817, 0.7349],\n",
    "                        [0.2432, 1.6161],\n",
    "                        [0.2194, 0.4353]])\n",
    "assert_allclose(distortion(ps), torch.FloatTensor([[0.1998, 0.2573],\n",
    "                                                   [1.2215, 0.9145],\n",
    "                                                   [0.4196, 2.1581],\n",
    "                                                   [0.2393, 0.4630]]), atol=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-05T23:35:34.332425Z",
     "start_time": "2020-05-05T23:35:34.329342Z"
    }
   },
   "source": [
    "### Wang08 distortion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.738799Z",
     "start_time": "2020-08-06T18:57:52.730417Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Wang08Distortion(nn.Module):\n",
    "    def __init__(self, d):\n",
    "        super().__init__()\n",
    "        self.d = tensors2parameters(d)\n",
    "    \n",
    "    def __repr__(self):\n",
    "        k1, k2, p, t = self.d\n",
    "        return f'{self.__class__.__name__}(k1:{k1:.4} k2:{k2:.4} p:{p:.4} t:{t:.4})'\n",
    "    \n",
    "    def forward(self, ps):\n",
    "        k1, k2, p, t = self.d\n",
    "        xs, ys = ps[:,0], ps[:,1]\n",
    "        \n",
    "        # Radial distortion\n",
    "        rs = xs**2 + ys**2\n",
    "        xs_r = xs*(1 + k1*rs + k2*rs**2)\n",
    "        ys_r = ys*(1 + k1*rs + k2*rs**2)\n",
    "\n",
    "        # Image plane (small angle approximation) rotation distortion\n",
    "        xs_d = xs_r/(-p*xs_r + t*ys_r + 1)\n",
    "        ys_d = ys_r/(-p*xs_r + t*ys_r + 1)\n",
    "        \n",
    "        return stackify((xs_d, ys_d), dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NOTE: possibly refactor the radial distortion part since it appears twice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.748493Z",
     "start_time": "2020-08-06T18:57:52.740380Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Wang08Distortion(k1:0.01 k2:0.02 p:0.03 t:0.04)"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "distortion = Wang08Distortion(torch.FloatTensor([.01, 0.02, 0.03, 0.04]))\n",
    "distortion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.754550Z",
     "start_time": "2020-08-06T18:57:52.750003Z"
    }
   },
   "outputs": [],
   "source": [
    "ps = torch.FloatTensor([[0.1900, 0.2467],\n",
    "                        [0.9817, 0.7349],\n",
    "                        [0.2432, 1.6161],\n",
    "                        [0.2194, 0.4353]])\n",
    "assert_allclose(distortion(ps), torch.FloatTensor([[0.1894, 0.2460],\n",
    "                                                   [1.0409, 0.7792],\n",
    "                                                   [0.2665, 1.7711],\n",
    "                                                   [0.2178, 0.4321]]), atol=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Camera matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-03T16:34:07.608445Z",
     "start_time": "2020-05-03T16:34:07.605173Z"
    }
   },
   "source": [
    "This performs the application of the Camera matrix. Implement the reconstruction in the `a2A()` method for child classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.759998Z",
     "start_time": "2020-08-06T18:57:52.755703Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def a2A(alpha_x, alpha_y, x_o, y_o):\n",
    "    zero, one = alpha_x.new_tensor(0), alpha_x.new_tensor(1)\n",
    "\n",
    "    return stackify(((alpha_x,    zero, x_o),\n",
    "                     (   zero, alpha_y, y_o),\n",
    "                     (   zero,    zero, one)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test to make sure reconstruction is correct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.766363Z",
     "start_time": "2020-08-06T18:57:52.760981Z"
    }
   },
   "outputs": [],
   "source": [
    "a = torch.FloatTensor([1.5, -2.5, 3.5])\n",
    "A = a2A(alpha_x=a[0], alpha_y=a[0], x_o=a[1], y_o=a[2])\n",
    "assert_allclose(A, torch.FloatTensor([[ 1.5,  0.0, -2.5],\n",
    "                                      [ 0.0,  1.5,  3.5],\n",
    "                                      [ 0.0,  0.0,  1.0]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.772477Z",
     "start_time": "2020-08-06T18:57:52.767350Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Cam(Inversible):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "    \n",
    "    def a2A(self):\n",
    "        raise NotImplementedError('Please implement camera matrix reconstruction')\n",
    "    \n",
    "    def forward_param(self): return self.a2A()\n",
    "    def inverse_param(self): return torch.inverse(self.forward_param()) # Overwrite this in child classes\n",
    "        \n",
    "    def forward(self, ps, inverse=False):\n",
    "        A = self.forward_param() if not inverse else self.inverse_param()\n",
    "        return pmm(ps, A, aug=True)\n",
    "    def inverse(self, ps): return self.forward(ps, inverse=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Single focal length"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-03T16:37:13.178310Z",
     "start_time": "2020-05-03T16:37:13.169302Z"
    }
   },
   "source": [
    "Assumes there is only a single `alpha`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.781183Z",
     "start_time": "2020-08-06T18:57:52.773904Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class CamSF(Cam):\n",
    "    def __init__(self, A):\n",
    "        super().__init__()\n",
    "        assert_allclose(A[0,0], A[1,1]) # Check that alpha_x and alpha_y are the same\n",
    "        self.a = tensors2parameters(A[[0,0,1],[0,2,2]])\n",
    "    \n",
    "    def __repr__(self):\n",
    "        alpha, x_o, y_o = self.a\n",
    "        return f'{self.__class__.__name__}(alpha:{alpha:.4} x_o:{x_o:.4} y_o:{y_o:.4})'\n",
    "    \n",
    "    def a2A(self):\n",
    "        alpha, x_o, y_o = self.a\n",
    "        return a2A(alpha_x=alpha, alpha_y=alpha, x_o=x_o, y_o=y_o)\n",
    "    \n",
    "    def inverse_param(self): # Overwrite parent method\n",
    "        alpha, x_o, y_o = self.a\n",
    "        return a2A(alpha_x=1/alpha, \n",
    "                   alpha_y=1/alpha, \n",
    "                   x_o=-x_o/alpha, \n",
    "                   y_o=-y_o/alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.788809Z",
     "start_time": "2020-08-06T18:57:52.782257Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CamSF(alpha:1.5 x_o:-2.5 y_o:3.5)"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = torch.FloatTensor([1.5, -2.5, 3.5])\n",
    "A = a2A(alpha_x=a[0], alpha_y=a[0], x_o=a[1], y_o=a[2])\n",
    "cam = CamSF(A)\n",
    "cam"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.795900Z",
     "start_time": "2020-08-06T18:57:52.790630Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.8154, 0.8791],\n",
       "        [0.6848, 0.1102],\n",
       "        [0.6551, 0.7518],\n",
       "        [0.6211, 0.1870]])"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ps = torch.rand(4,2)\n",
    "ps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.801500Z",
     "start_time": "2020-08-06T18:57:52.796885Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_inversible(cam, ps, pmm(ps, A, aug=True), atol=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:52.807658Z",
     "start_time": "2020-08-06T18:57:52.802655Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose(cam.get_param(), A)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-03T20:36:21.604772Z",
     "start_time": "2020-05-03T20:36:21.600609Z"
    }
   },
   "source": [
    "# Build"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:57:53.082955Z",
     "start_time": "2020-08-06T18:57:52.808737Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "IPython.notebook.save_checkpoint();"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converted modules.ipynb.\n"
     ]
    }
   ],
   "source": [
    "build_notebook()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  },
  "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": 2
}