fit_ground_plane.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plane Fitting and Keyframe Projection\n",
    "\n",
    "The following script has two main tasks:\n",
    "1. Fit a plane into the 3D map points which were triangulated by the SLAM system\n",
    "2. Project the individual key frames onto this plane and blend them into a single large mosaic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib ipympl\n",
    "import pickle\n",
    "import numpy as np\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from pytransform3d.rotations import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_points = pickle.load(open(\"map_points.pkl\", \"rb\"))\n",
    "kf_visible_map_points = pickle.load(open(\"kf_visible_map_points.pkl\", \"rb\"))\n",
    "kf_poses = pickle.load(open(\"kf_poses.pkl\", \"rb\"))\n",
    "kf_frames = pickle.load(open(\"kf_frames.pkl\", \"rb\"))\n",
    "kf_kp_matched = pickle.load(open(\"kf_kp_matched.pkl\", \"rb\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fit a Plane into the map points using a RANSAC scheme"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import numpy.linalg as la\n",
    "\n",
    "eps = 0.00001\n",
    "\n",
    "def svd(A):\n",
    "    u, s, vh = la.svd(A)\n",
    "    S = np.zeros(A.shape)\n",
    "    S[:s.shape[0], :s.shape[0]] = np.diag(s)\n",
    "    return u, S, vh\n",
    "\n",
    "\n",
    "def fit_plane_LSE(points):\n",
    "    # points: Nx4 homogeneous 3d points\n",
    "    # return: 1d array of four elements [a, b, c, d] of\n",
    "    # ax+by+cz+d = 0\n",
    "    assert points.shape[0] >= 3 # at least 3 points needed\n",
    "    U, S, Vt = svd(points)\n",
    "    null_space = Vt[-1, :]\n",
    "    return null_space\n",
    "\n",
    "\n",
    "def get_point_dist(points, plane):\n",
    "    # return: 1d array of size N (number of points)\n",
    "    dists = np.abs(points @ plane) / np.sqrt(plane[0]**2 + plane[1]**2 + plane[2]**2)\n",
    "    return dists\n",
    "\n",
    "\n",
    "def fit_plane_LSE_RANSAC(points, iters=1000, inlier_thresh=0.05, num_support_points=None, return_outlier_list=False):\n",
    "    # points: Nx4 homogeneous 3d points\n",
    "    # num_support_points: If None perform LSE fit with all points, else pick `num_support_points` random points for fitting\n",
    "    # return: \n",
    "    #   plane: 1d array of four elements [a, b, c, d] of ax+by+cz+d = 0\n",
    "    #   inlier_list: 1d array of size N of inlier points\n",
    "    max_inlier_num = -1\n",
    "    max_inlier_list = None\n",
    "    \n",
    "    N = points.shape[0]\n",
    "    assert N >= 3\n",
    "\n",
    "    for i in range(iters):\n",
    "        chose_id = np.random.choice(N, 3, replace=False)\n",
    "        chose_points = points[chose_id, :]\n",
    "        tmp_plane = fit_plane_LSE(chose_points)\n",
    "        \n",
    "        dists = get_point_dist(points, tmp_plane)\n",
    "        tmp_inlier_list = np.where(dists < inlier_thresh)[0]\n",
    "        tmp_inliers = points[tmp_inlier_list, :]\n",
    "        num_inliers = tmp_inliers.shape[0]\n",
    "        if num_inliers > max_inlier_num:\n",
    "            max_inlier_num = num_inliers\n",
    "            max_inlier_list = tmp_inlier_list\n",
    "        \n",
    "        #print('iter %d, %d inliers' % (i, max_inlier_num))\n",
    "\n",
    "    final_points = points[max_inlier_list, :]\n",
    "    if num_support_points:\n",
    "        max_support_points = np.min((num_support_points, final_points.shape[0]))\n",
    "        support_idx = np.random.choice(np.arange(0, final_points.shape[0]), max_support_points, replace=False)\n",
    "        support_points = final_points[support_idx, :]\n",
    "    else:\n",
    "        support_points = final_points\n",
    "    print(final_points.shape)\n",
    "    plane = fit_plane_LSE(support_points)\n",
    "    \n",
    "    fit_variance = np.var(get_point_dist(final_points, plane))\n",
    "    print('RANSAC fit variance: %f' % fit_variance)\n",
    "    print(plane)\n",
    "\n",
    "    dists = get_point_dist(points, plane)\n",
    "\n",
    "    select_thresh = inlier_thresh * 1\n",
    "\n",
    "    inlier_list = np.where(dists < select_thresh)[0]\n",
    "    if not return_outlier_list:\n",
    "        return plane, inlier_list\n",
    "    else:\n",
    "        outlier_list = np.where(dists >= select_thresh)[0]\n",
    "        return plane, inlier_list, outlier_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_points_h = cv2.convertPointsToHomogeneous(map_points).reshape(-1, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(29721, 4)\n",
      "RANSAC fit variance: 0.046220\n",
      "[ 4.51320124e-04 -7.53968813e-05  4.98553894e-02 -9.98756342e-01]\n",
      "(29232,) 1.1003984468989074\n"
     ]
    }
   ],
   "source": [
    "# limiting num_support_points speeds up this operation significantly\n",
    "t0 = time.perf_counter()\n",
    "plane, inlier_list = fit_plane_LSE_RANSAC(map_points_h, iters=1000, inlier_thresh=1, num_support_points=3000, return_outlier_list=False)\n",
    "dt = time.perf_counter() - t0\n",
    "print(inlier_list.shape, dt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Invert plane if normal points in the \"wrong\" direction\n",
    "\n",
    "During fitting the plane can be either fitted in a way that the normal points towards the world origin (parameter d > 0) or in the opposite direction (d < 0). The following calculations assume the plane normal to point into the direction of the world origin. So, in case the plane is fitted with d < 0, the plane normal has to be inverted, that is \"flipped\" by 180° along its own axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [],
   "source": [
    "if plane[3] < 0:\n",
    "    plane *= -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-4.51320124e-04  7.53968813e-05 -4.98553894e-02  9.98756342e-01]\n"
     ]
    }
   ],
   "source": [
    "print(plane)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an orthonormal plane base "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plane_to_hessian(plane):\n",
    "    \"\"\"Convert plane to Hessian normal form (n.x + p = 0)\"\"\"\n",
    "    a, b, c, d = plane\n",
    "    nn = np.sqrt(a**2+b**2+c**2)\n",
    "    n = np.array([a/nn, b/nn, c/nn])\n",
    "    p = d/nn\n",
    "    return n, p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_plane_base(plane):\n",
    "    \"\"\"Create a orthonormal plane coordinate base.\n",
    "    \n",
    "    Args:\n",
    "    \n",
    "        plane (`numpy.ndarray`): Shape (4,). Plane coefficients [a, b, c, d] which fullfill\n",
    "            ax + by + cz + d = 0.\n",
    "    \n",
    "    Returns:\n",
    "    \n",
    "        base (`numpy.ndarray`) Shape (3, 4). right-handend orthonormal base in which 2D plane points are \n",
    "            expressed. Column vectors of the array correspond to the origin point O of the base, the first and \n",
    "            second base vector u, v and the third base vector n which is the plane normal.\n",
    "    \"\"\"\n",
    "    # get two points on the plane for x = 0, y = 0 and x = 1, y = 0, respectively\n",
    "    z1 = -1*plane[3]/plane[2]\n",
    "    p1 = np.array([0, 0, z1])\n",
    "    z2 = -1*(plane[0]+plane[3])/plane[2]\n",
    "    p2 = np.array([1, 0, z2])\n",
    "    # first plane base vector u\n",
    "    u = (p2-p1)/np.linalg.norm(p2-p1)\n",
    "    # third plane base vector := plane normal\n",
    "    n, _ = plane_to_hessian(plane)\n",
    "    # second plane base vector v \n",
    "    uxn = np.cross(u, n)\n",
    "    v = uxn/np.linalg.norm(uxn)\n",
    "    # assemble base as single matrix\n",
    "    base = np.stack([p1.T, u.T, v.T, n.T], axis=1)\n",
    "    return base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.00000000e+00  9.99959028e-01  1.36891906e-05 -9.05220323e-03]\n",
      " [ 0.00000000e+00  0.00000000e+00  9.99998857e-01  1.51224786e-03]\n",
      " [ 2.00330667e+01 -9.05221358e-03  1.51218590e-03 -9.99957884e-01]]\n"
     ]
    }
   ],
   "source": [
    "base = make_plane_base(plane)\n",
    "print(base)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert camera poses from world to plane coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_twist(R, t):\n",
    "    \"\"\"Convert a 3x3 rotation matrix and translation vector (shape (3,))\n",
    "    into a 6D twist coordinate (shape (6,)).\"\"\"\n",
    "    r, _ = cv2.Rodrigues(R)\n",
    "    twist = np.zeros((6,))\n",
    "    twist[:3] = r.reshape(3,)\n",
    "    twist[3:] = t.reshape(3,)\n",
    "    return twist\n",
    "\n",
    "def from_twist(twist):\n",
    "    \"\"\"Convert a 6D twist coordinate (shape (6,)) into a 3x3 rotation matrix\n",
    "    and translation vector (shape (3,)).\"\"\"\n",
    "    r = twist[:3].reshape(3, 1)\n",
    "    t = twist[3:].reshape(3, 1)\n",
    "    R, _ = cv2.Rodrigues(r)\n",
    "    return R, t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first compute R_plane, t_plane from given coordinate systems\n",
    "# see: https://math.stackexchange.com/questions/1125203/finding-rotation-axis-and-angle-to-align-two-3d-vector-bases\n",
    "def get_plane_pose(plane_O, plane_u, plane_v, plane_n):\n",
    "    \"\"\"Compute the plane pose R_plane, t_plane in world coordinates.\n",
    "    \n",
    "    Assumes the world coordinate system to have rotation R = I and zero translation.\n",
    "    \n",
    "    Args:\n",
    "        plane_O (`numpy.ndarray`): Shape (3,). Origin of the plane coordinate base.\n",
    "        \n",
    "        plane_u (`numpy.ndarray`): Shape (3,). First base vector of the plane coordinate system.\n",
    "        \n",
    "        plane_v (`numpy.ndarray`): Shape (3,). Second base vector of the plane coordinate system.\n",
    "        \n",
    "        plane_n (`numpy.ndarray`): Shape (3,). Third base vector of the plane coordinate system, \n",
    "            corresponds to plane normal.\n",
    "            \n",
    "    Returns:\n",
    "        R_plane (`numpy.ndarray`), t_plane (`numpy.ndarray`): Rotation matrix with shape (3, 3) and \n",
    "        translation vector with shape (3,) which describe the pose of the plane coordinate base in\n",
    "        the world coordinate frame.\n",
    "    \"\"\"\n",
    "    t_plane = plane_O\n",
    "    world_x = np.array([1, 0, 0])\n",
    "    world_y = np.array([0, 1, 0])\n",
    "    world_z = np.array([0, 0, 1])\n",
    "    R_plane = np.outer(plane_u, world_x) + np.outer(plane_v, world_y) + np.outer(plane_n, world_z)\n",
    "    return R_plane, t_plane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 9.99959028e-01  1.36891906e-05 -9.05220323e-03]\n",
      " [ 0.00000000e+00  9.99998857e-01  1.51224786e-03]\n",
      " [-9.05221358e-03  1.51218590e-03 -9.99957884e-01]] [ 0.          0.         20.03306673]\n"
     ]
    }
   ],
   "source": [
    "R_plane, t_plane = get_plane_pose(base[:, 0], base[:, 1], base[:, 2], base[:, 3])\n",
    "print(R_plane, t_plane)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pose_world_to_plane(R_plane, t_plane, R_pose, t_pose):\n",
    "    \"\"\"Map keyframe poses from the world to plane coordinate frame.\n",
    "\n",
    "    Args:\n",
    "        R_plane (`numpy.ndarray`): Shape (3, 3). Rotation matrix of plane coordinate system in world coordinate frame.\n",
    "        t_plane (`numpy.ndarray`): Shape (3,). Translation vector of plane coordinate system in world coordinate frame.\n",
    "        \n",
    "        R_pose (`numpy.ndarray`): Shape (3, 3). Rotation matrix of keyframe in world coordinate frame.\n",
    "        t_pose (`numpy.ndarray`): Shape (3,). Translation vector of keyframe in world coordinate frame.\n",
    "\n",
    "    Returns:\n",
    "        R_pose_plane (`numpy.ndarray`), t_pose_plane (`numpy.ndarray`): Rotation matrix with shape (3, 3) and \n",
    "        translation vector with shape (3,) of the keyframe in plane coordinate frame.\n",
    "    \"\"\"\n",
    "    t_pose_plane = t_plane.reshape(3,) + R_plane @ t_pose.reshape(3,)\n",
    "    R_pose_plane = R_plane @ R_pose\n",
    "    return R_pose_plane, t_pose_plane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert all keyframe poses to plane coordinates\n",
    "kf_poses_plane = []\n",
    "for pose in kf_poses:\n",
    "    R_pose, t_pose = from_twist(pose)\n",
    "    R_kf_plane, t_kf_plane = pose_world_to_plane(R_plane, t_plane, R_pose, t_pose)\n",
    "    kf_poses_plane.append((R_kf_plane, t_kf_plane))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "522c695f3d224b59bad4936c2cd4d2f8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FigureCanvasNbAgg()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "R_world = np.eye(3)\n",
    "t_world = np.zeros((3,))\n",
    "\n",
    "fig = plt.figure(figsize=(7, 7))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "plot_basis(ax, R_world, t_world)\n",
    "plot_basis(ax, R_plane, t_plane)\n",
    "\n",
    "R_0, t_0 = from_twist(kf_poses[0])\n",
    "plot_basis(ax, R_0, t_0.reshape(3,))\n",
    "R_1, t_1 = from_twist(kf_poses[1])\n",
    "plot_basis(ax, R_1, t_1.reshape(3,))\n",
    "\n",
    "R_0, t_0 = kf_poses_plane[0]\n",
    "plot_basis(ax, R_0, t_0.reshape(3,))\n",
    "R_1, t_1 = kf_poses_plane[1]\n",
    "plot_basis(ax, R_1, t_1.reshape(3,))\n",
    "\n",
    "ax.scatter(base[0, 0], base[1, 0], base[2, 0], c=\"orange\")\n",
    "ax.scatter(base[0, 0]+base[0, 1], base[1, 0]+base[1, 1], base[2, 0]+base[2, 1], c=\"red\")\n",
    "ax.scatter(base[0, 0]+base[0, 2], base[1, 0]+base[1, 2], base[2, 0]+base[2, 2], c=\"green\") \n",
    "ax.scatter(base[0, 0]+base[0, 3], base[1, 0]+base[1, 3], base[2, 0]+base[2, 3], c=\"blue\")\n",
    "ax.set_xlim([-20, 20])\n",
    "ax.set_ylim([-20, 40])\n",
    "ax.set_zlim([0, 40])\n",
    "ax.set_aspect(1.0)\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"y\")\n",
    "ax.set_zlabel(\"z\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Project image corners onto plane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {},
   "outputs": [],
   "source": [
    "w = 1920\n",
    "h = 1080\n",
    "fx = 1184.51770\n",
    "fy = 1183.63810\n",
    "cx = 978.30778\n",
    "cy = 533.85598"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {},
   "outputs": [],
   "source": [
    "# corners of the image\n",
    "img_points = np.array([[0, 0],\n",
    "                       [w, 0],\n",
    "                       [0, h],\n",
    "                       [w, h]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [],
   "source": [
    "def unproject(img_points, fx, fy, cx, cy):\n",
    "    \"\"\"Unproject image points with shape (-1, 2) to camera coordinates.\"\"\"\n",
    "    camera_points = np.zeros((img_points.shape[0], 3))\n",
    "    camera_points[:, 0] = (img_points[:, 0]-cx)/fx\n",
    "    camera_points[:, 1] = (img_points[:, 1]-cy)/fy\n",
    "    camera_points[:, 2] = 1.0\n",
    "    return camera_points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-0.82591234, -0.45102974,  1.        ],\n",
       "       [ 0.79500055, -0.45102974,  1.        ],\n",
       "       [-0.82591234,  0.46141132,  1.        ],\n",
       "       [ 0.79500055,  0.46141132,  1.        ]])"
      ]
     },
     "execution_count": 146,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "camera_points = unproject(img_points, fx, fy, cx, cy)\n",
    "camera_points  # image corners in camera coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {},
   "outputs": [],
   "source": [
    "# adapted from: https://github.com/zdzhaoyong/Map2DFusion/blob/master/src/Map2DRender.cpp\n",
    "warped_frames = [None for _ in range(len(kf_frames))]\n",
    "warped_masks = [None for _ in range(len(kf_frames))]\n",
    "length_pixel = 0.01\n",
    "view_min = np.array([1e6, 1e6])\n",
    "view_max = np.array([-1e6, -1e6])\n",
    "\n",
    "sizes = np.zeros((len(kf_frames), 2), dtype=np.int)\n",
    "corners_world = np.zeros((len(kf_frames), 2))\n",
    "\n",
    "for idx, (pose, frame) in enumerate(zip(kf_poses_plane, kf_frames)):\n",
    "    \n",
    "    frame_okay = True\n",
    "    cur_view_min = np.array([1e6, 1e6])\n",
    "    cur_view_max = np.array([-1e6, -1e6])\n",
    "    \n",
    "    # generally applicable because the world coordinate base coincides with the first keyframe and plane \n",
    "    # points will always have a positive z_world coordinate. Further, we ensure the plane normal to point\n",
    "    # towards the world origin.\n",
    "    downlook = np.array([0,0,-1])\n",
    "    \n",
    "    # project image corners from camera to plane coordinates\n",
    "    plane_points = np.zeros((camera_points.shape[0], 2))\n",
    "    for i, camera_point in enumerate(camera_points):\n",
    "        axis = pose[0] @ camera_point\n",
    "        if np.dot(axis, downlook) < 0.4:\n",
    "            print(\"Camera axis is deviating too much from 'down' direction. Skipping to next keyframe.\")\n",
    "            frame_okay = False\n",
    "            break\n",
    "        axis = pose[1] - axis*(pose[1][-1]/axis[-1])\n",
    "        plane_points[i, :] = axis[:2]\n",
    "        \n",
    "    if not frame_okay:\n",
    "        continue\n",
    "    \n",
    "    # expand viewport of current frame\n",
    "    for i, plane_point in enumerate(plane_points):\n",
    "        if plane_point[0] < cur_view_min[0]:\n",
    "            cur_view_min[0] = plane_point[0]\n",
    "        if plane_point[1] < cur_view_min[1]:\n",
    "            cur_view_min[1] = plane_point[1]\n",
    "        if plane_point[0] > cur_view_max[0]:\n",
    "            cur_view_max[0] = plane_point[0]\n",
    "        if plane_point[1] > cur_view_max[1]:\n",
    "            cur_view_max[1] = plane_point[1]\n",
    "            \n",
    "    # expand overall viewport if necessary\n",
    "    if cur_view_min[0] < view_min[0]:\n",
    "        view_min[0] = cur_view_min[0]\n",
    "    if cur_view_min[1] < view_min[1]:\n",
    "        view_min[1] = cur_view_min[1]\n",
    "    if cur_view_max[0] > view_max[0]:\n",
    "        view_max[0] = cur_view_max[0]\n",
    "    if cur_view_max[1] > view_max[1]:\n",
    "        view_max[1] = cur_view_max[1]\n",
    "\n",
    "    corners_world[idx, :] = cur_view_min\n",
    "    sizes[idx, :] = ((cur_view_max - cur_view_min)/length_pixel)    \n",
    "    dst_points = (plane_points - cur_view_min)/length_pixel\n",
    "    \n",
    "    # find homography between camera and ground plane points\n",
    "    transmtx = cv2.getPerspectiveTransform(img_points.astype(np.float32), dst_points.astype(np.float32))\n",
    "    \n",
    "    mask = np.full(frame.shape[0:2], 255, dtype=np.uint8)\n",
    "    \n",
    "    # warp image & mask\n",
    "    warped_frame = cv2.warpPerspective(frame, transmtx, tuple(sizes[idx, :]), cv2.INTER_CUBIC, cv2.BORDER_REFLECT)\n",
    "    warped_mask = cv2.warpPerspective(mask, transmtx, tuple(sizes[idx, :]), cv2.INTER_NEAREST)\n",
    "    warped_frames[idx] = warped_frame\n",
    "    warped_masks[idx] = warped_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "664ed9da28494435884e80d966c75366",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FigureCanvasNbAgg()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "f, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))\n",
    "ax1.imshow(warped_frames[0][:, :, ::-1])\n",
    "ax2.imshow(warped_masks[0][:, :])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cb8c0e3057374c388c7f5a952e938f3d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FigureCanvasNbAgg()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "R_world = np.eye(3)\n",
    "t_world = np.zeros((3,))\n",
    "\n",
    "fig = plt.figure(figsize=(7, 7))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "plot_basis(ax, R_world, t_world)\n",
    "plot_basis(ax, R_plane, t_plane)\n",
    "\n",
    "R_0, t_0 = from_twist(kf_poses[0])\n",
    "plot_basis(ax, R_0, t_0.reshape(3,))\n",
    "R_1, t_1 = from_twist(kf_poses[1])\n",
    "plot_basis(ax, R_1, t_1.reshape(3,))\n",
    "\n",
    "R_0, t_0 = kf_poses_plane[0]\n",
    "plot_basis(ax, R_0, t_0.reshape(3,))\n",
    "R_1, t_1 = kf_poses_plane[1]\n",
    "plot_basis(ax, R_1, t_1.reshape(3,))\n",
    "\n",
    "#ax.scatter(points_world[:1000, 0], points_world[:1000, 1], points_world[:1000, 2], s=1, c=\"red\")\n",
    "#ax.scatter(points_plane[:1000, 0], points_plane[:1000, 1], s=1, c=\"green\")\n",
    "\n",
    "ax.scatter(camera_points[:, 0], camera_points[:, 1], camera_points[:, 2], c=\"cyan\")\n",
    "ax.scatter(plane_points[:, 0], plane_points[:, 1], c=\"magenta\")\n",
    "\n",
    "ax.scatter(base[0, 0], base[1, 0], base[2, 0], c=\"orange\")\n",
    "ax.scatter(base[0, 0]+base[0, 1], base[1, 0]+base[1, 1], base[2, 0]+base[2, 1], c=\"red\")\n",
    "ax.scatter(base[0, 0]+base[0, 2], base[1, 0]+base[1, 2], base[2, 0]+base[2, 2], c=\"green\") \n",
    "ax.scatter(base[0, 0]+base[0, 3], base[1, 0]+base[1, 3], base[2, 0]+base[2, 3], c=\"blue\")\n",
    "ax.set_xlim([-20, 20])\n",
    "ax.set_ylim([-20, 40])\n",
    "ax.set_zlim([0, 40])\n",
    "ax.set_aspect(1.0)\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"y\")\n",
    "ax.set_zlabel(\"z\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stitch frames together using OpenCV Blenders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "metadata": {},
   "outputs": [],
   "source": [
    "corners_images = (corners_world - view_min)/length_pixel\n",
    "corners_images = corners_images.astype(np.int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "result ROI: [   0    0 4725 7845]\n"
     ]
    }
   ],
   "source": [
    "# define ROI of stitched mosaic\n",
    "bottom_right = np.max(corners_images, axis=0) + np.array([sizes[np.argmax(corners_images, axis=0)[0], 0], sizes[np.argmax(corners_images, axis=0)[1], 1]])\n",
    "result_roi = np.array([0, 0, bottom_right[0], bottom_right[1]])\n",
    "print(\"result ROI:\", result_roi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "304.4159366721789\n"
     ]
    }
   ],
   "source": [
    "# compute number of channels\n",
    "blend_strength = 5\n",
    "result_roi_area = (result_roi[2] - result_roi[0]) * (result_roi[3] - result_roi[1])\n",
    "blend_width = np.sqrt(result_roi_area) * blend_strength/100\n",
    "print(blend_width)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Multiband Blender"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using bands: 8\n"
     ]
    }
   ],
   "source": [
    "num_bands = int(np.ceil(np.log(blend_width)/np.log(2)) - 1)\n",
    "print(\"Using bands:\", num_bands)\n",
    "\n",
    "blender = cv2.detail_MultiBandBlender(try_gpu=False, num_bands=num_bands)\n",
    "blender.prepare(result_roi)\n",
    "\n",
    "for idx, (frame, mask) in enumerate(zip(warped_frames, warped_masks)):\n",
    "    if frame is not None:\n",
    "        blender.feed(frame, mask, tuple(corners_images[idx]))\n",
    "    \n",
    "result, result_mask = blender.blend(None, None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5e78d19a18f44eda811f8e24dd00e3a2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FigureCanvasNbAgg()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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"
     ]
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(10, 6))\n",
    "ax = fig.add_subplot(111)\n",
    "ax.imshow(result[:, :, ::-1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 155,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cv2.imwrite(\"result.jpg\", result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Feather Blender \n",
    "(result looks better than with multiband blender)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "metadata": {},
   "outputs": [],
   "source": [
    "blender = cv2.detail_FeatherBlender(sharpness=1./blend_width)\n",
    "blender.prepare(result_roi)\n",
    "\n",
    "for idx, (frame, mask) in enumerate(zip(warped_frames, warped_masks)):\n",
    "    if frame is not None:\n",
    "        blender.feed(frame.astype(np.int16), mask, tuple(corners_images[idx]))\n",
    "    \n",
    "result, result_mask = blender.blend(None, None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b2b8973e9ede465dafb41a29e3aa6e38",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FigureCanvasNbAgg()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(10, 6))\n",
    "ax = fig.add_subplot(111)\n",
    "ax.imshow(result[:, :, ::-1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "66c4cecb252b4ded90e4b91e71cb51be",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FigureCanvasNbAgg()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(10, 6))\n",
    "ax = fig.add_subplot(111)\n",
    "ax.imshow(result_mask)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 142,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cv2.imwrite(\"result_feather.jpg\", result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TODO:\n",
    "- realize the same blending as in Map2DFusion. The raw cells below contain some code that creates a weighted mask which is to be used for blending. Note, that the MultiBand blender does not work well with this weighted mask."
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "weight_image = np.full((h, w), 255, dtype=np.uint8)\n",
    "\n",
    "center = np.array([0.5*w, 0.5*h])\n",
    "dismaxinv = 1./np.linalg.norm(center)\n",
    "\n",
    "for i in range(0, h):\n",
    "    for j in range(0, w):\n",
    "        \n",
    "        dis = (i - center[1])**2 + (j - center[0])**2\n",
    "        dis = 1 - np.sqrt(dis) * dismaxinv\n",
    "        \n",
    "        weight_image[i, j] = dis**2 * 254\n",
    "        if weight_image[i, j] < 1: \n",
    "            weight_image[i, j] = 1"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "fig = plt.figure(figsize=(10, 6))\n",
    "ax = fig.add_subplot(111)\n",
    "weight_image = cv2.cvtColor(weight_image, cv2.COLOR_GRAY2BGR)\n",
    "ax.imshow(weight_image[:, :, ::-1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}