{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "import open3d as o3d\n",
    "import open3d.core as o3c\n",
    "\n",
    "if o3d.__DEVICE_API__ == 'cuda':\n",
    "    import open3d.cuda.pybind.t.pipelines.registration as treg\n",
    "else:\n",
    "    import open3d.cpu.pybind.t.pipelines.registration as treg\n",
    "\n",
    "import numpy as np\n",
    "import sys\n",
    "import os\n",
    "import time\n",
    "\n",
    "# monkey patches visualization and provides helpers to load geometries\n",
    "sys.path.append('..')\n",
    "import open3d_tutorial as o3dtut\n",
    "# change to True if you want to interact with the visualization windows\n",
    "o3dtut.interactive = not \"CI\" in os.environ"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Robust Kernel\n",
    "\n",
    "This tutorial demonstrates the use of robust kernels in the context of outlier rejection. For this particular tutorial, we will be using the ICP (Iterative Closest Point) registration algorithm as the _target_ problem where we want to deal with outliers. Even so, the theory applies to any given optimization problem and not just for ICP. For the moment the robust kernels have been only implemented for the `PointToPlane` ICP.\n",
    "\n",
    "The notation and some of the kernels implemented in `Open3D` has been inspired by the publication \"Analysis of Robust Functions for Registration Algorithms\" [\\[Babin2019\\]](../reference.html#babin2019)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    " \n",
    "**Note:** \n",
    "\n",
    "This tutorial and the original implementation of the robust kernels in `Open3D` was contributed by **Ignacio Vizzo** and **Cyrill Stachniss** from the University of Bonn.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Point-to-plane ICP using Robust Kernels\n",
    "\n",
    "The standard point-to-plane ICP algorithm [\\[ChenAndMedioni1992\\]](../reference.html#chenandmedioni1992) minimizes this objective function:\n",
    "\n",
    "\\begin{equation}\n",
    "E(\\mathbf{T}) = \\sum_{(\\mathbf{p},\\mathbf{q})\\in\\mathcal{K}}\\big((\\mathbf{p} - \\mathbf{T}\\mathbf{q})\\cdot\\mathbf{n}_{\\mathbf{p}}\\big)^{2},\n",
    "\\end{equation}\n",
    "\n",
    "where $\\mathbf{n}_{\\mathbf{p}}$ is the normal of point $\\mathbf{p}$ and $\\mathcal{K}$ is the correspondence set between the target point cloud $\\mathbf{P}$, and source point cloud $\\mathbf{Q}$.\n",
    "\n",
    "If we call $r_i(\\mathbf{T})$ is $i^\\mathrm{th}$ the residual, for a given pair of correspondences $(\\mathbf{p},\\mathbf{q})\\in\\mathcal{K}$ we can rewrite the objective function as.\n",
    "\n",
    "\\begin{equation}\n",
    "E(\\mathbf{T}) = \\sum_{(\\mathbf{p},\\mathbf{q})\\in\\mathcal{K}}\\big((\\mathbf{p} - \\mathbf{T}\\mathbf{q})\\cdot\\mathbf{n}_{\\mathbf{p}}\\big)^{2} = \\sum_{i=1}^{N} \\big({r_i(\\mathbf{T})}\\big)^2\n",
    "\\end{equation}\n",
    "\n",
    "The optimization problem above can also be solved by using the iteratively reweighted least-squares (IRLS) approach, which solves a sequence of weighted least squares problems:\n",
    "\n",
    "\\begin{equation}\n",
    "E(\\mathbf{T}) =  \\sum_{i=1}^{N} w_i \\big({r_i(\\mathbf{T})}\\big)^2\n",
    "\\end{equation}\n",
    "\n",
    "### Outlier Rejection with Robust Kernels\n",
    "\n",
    "The main idea of a robust loss is to downweight large residuals that are assumed to be caused from outliers such that their influence on the solution is reduced. This is achieved by optimizing $E(\\mathbf{T})$ as:\n",
    "\n",
    "\\begin{equation}\n",
    "E(\\mathbf{T}) = \\sum_{(\\mathbf{p},\\mathbf{q})\\in\\mathcal{K}}\\rho\\big((\\mathbf{p} - \\mathbf{T}\\mathbf{q})\\cdot\\mathbf{n}_{\\mathbf{p}}\\big) = \\sum_{i=1}^{N} \\rho\\big({r_i(\\mathbf{T})}\\big),\n",
    "\\end{equation}\n",
    "\n",
    "where $\\rho(r)$ is also called the robust loss or kernel.\n",
    "\n",
    "We can see that exists a relation between the optimization formulation in IRLS and the one that uses the robust loss function. By setting the weight $w_i= \\frac{1}{r_i(\\mathbf{T})}\\rho'(r_i(\\mathbf{T}))$, we can solve the robust loss optimization problem using the existing techniques for weighted least-squares. Therefore, we can minimize the objective function using Gauss-Newton and determine increments by iteratively solving:\n",
    "\n",
    "\\begin{align}\n",
    "\\left(\\mathbf{J}^\\top \\mathbf{W} \\mathbf{J}\\right)^{-1}\\mathbf{J}^\\top\\mathbf{W}\\vec{r},\n",
    "\\end{align}\n",
    "\n",
    "where $\\mathbf{W} \\in \\mathbb{R}^{N\\times N}$ is a diagonal matrix containing weights $w_i$ for each residual $r_i$\n",
    "\n",
    "### How to use Robust Kernels in Open3D\n",
    "\n",
    "`icp`, `multi_scale_icp` can be called with a parameter `TransformationEstimationPointToPlane(loss)` or `TransformationEstimationForColoredICP(loss)`. Where `loss` is a given loss function(also called robust kernel). \n",
    "\n",
    "Internally, `TransormationEstimationPointToPlane(loss)` implements a function to compute the weighted residuals and Jacobian matrices of the point-to-plane ICP objective according to the provided robust kernel."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vanilla ICP vs Robust ICP\n",
    "\n",
    "To better show the advantages of using robust kernels in the registration, we add some artificially generated gaussian noise to the source point cloud."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Overriding visualization function, according to best camera view for colored-icp sample data.\n",
    "def draw_registration_result(source, target, transformation):\n",
    "    source_temp = source.clone()\n",
    "    target_temp = target.clone()\n",
    "\n",
    "    source_temp.transform(transformation)\n",
    "\n",
    "    # This is patched version for tutorial rendering.\n",
    "    # Use `draw` function for you application.\n",
    "    o3d.visualization.draw_geometries(\n",
    "        [source_temp.to_legacy(),\n",
    "         target_temp.to_legacy()],\n",
    "        zoom=0.4459,\n",
    "        front=[0.9288, -0.2951, -0.2242],\n",
    "        lookat=[1.6784, 2.0612, 1.4451],\n",
    "        up=[-0.3402, -0.9189, -0.1996])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "demo_icp_pcds = o3d.data.DemoICPPointClouds()\n",
    "source = o3d.t.io.read_point_cloud(demo_icp_pcds.paths[0])\n",
    "target = o3d.t.io.read_point_cloud(demo_icp_pcds.paths[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_noise(pcd, mu, sigma):\n",
    "    device = pcd.point[\"positions\"].device\n",
    "    noisy_pcd = pcd.cpu().clone()\n",
    "    noisy_pcd.point[\"positions\"] += o3d.core.Tensor(\n",
    "        np.random.normal(mu, sigma, size=noisy_pcd.point[\"positions\"].shape),\n",
    "        noisy_pcd.point[\"positions\"].dtype)\n",
    "    return noisy_pcd\n",
    "\n",
    "\n",
    "mu, sigma = 0, 0.1  # mean and standard deviation\n",
    "source_noisy = apply_noise(source, mu, sigma)\n",
    "\n",
    "print(\"Source PointCloud + noise:\")\n",
    "o3d.visualization.draw_geometries([source_noisy.to_legacy()],\n",
    "                                  zoom=0.4459,\n",
    "                                  front=[0.353, -0.469, -0.809],\n",
    "                                  lookat=[2.343, 2.217, 1.809],\n",
    "                                  up=[-0.097, -0.879, 0.467])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vanilla ICP\n",
    "\n",
    "We inspect how the results will look if we use the exact same parameters that were used in the above examples of ICP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initial alignment or source to target transform.\n",
    "init_source_to_target = np.asarray([[0.862, 0.011, -0.507, 0.5],\n",
    "                                    [-0.139, 0.967, -0.215, 0.7],\n",
    "                                    [0.487, 0.255, 0.835, -1.4],\n",
    "                                    [0.0, 0.0, 0.0, 1.0]])\n",
    "\n",
    "criteria = treg.ICPConvergenceCriteria(relative_fitness=0.000001,\n",
    "                                       relative_rmse=0.000001,\n",
    "                                       max_iteration=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimation = treg.TransformationEstimationPointToPlane()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_correspondence_distance = 0.02"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Vanilla point-to-plane ICP, max_correspondence_distance={}:\".format(\n",
    "    max_correspondence_distance))\n",
    "s = time.time()\n",
    "\n",
    "reg_point_to_plane = treg.icp(source_noisy, target, max_correspondence_distance,\n",
    "                              init_source_to_target, estimation)\n",
    "\n",
    "icp_time = time.time() - s\n",
    "\n",
    "print(\"Time taken by Point-To-Plane ICP: \", icp_time)\n",
    "print(\"Fitness: \", reg_point_to_plane.fitness)\n",
    "print(\"Inlier RMSE: \", reg_point_to_plane.inlier_rmse)\n",
    "\n",
    "draw_registration_result(source, target, reg_point_to_plane.transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tuning Vanilla ICP\n",
    "\n",
    "Given the fact that we are now dealing with gaussian noise, we might try to increase the threshold to search for nearest neighbors with the aim of improving the registration result.\n",
    "\n",
    "We can see that under these conditions and without a robust kernel, the traditional ICP has no chance to deal with  outliers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimation = treg.TransformationEstimationPointToPlane()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_correspondence_distance = 0.5\n",
    "\n",
    "voxel_size = 0.02"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Vanilla point-to-plane ICP, max_correspondence_distance={}:\".format(\n",
    "    max_correspondence_distance))\n",
    "\n",
    "s = time.time()\n",
    "\n",
    "reg_point_to_plane = treg.icp(source, target, max_correspondence_distance,\n",
    "                              init_source_to_target, estimation, criteria,\n",
    "                              voxel_size)\n",
    "\n",
    "icp_time = time.time() - s\n",
    "\n",
    "print(\"Time taken by Point-To-Plane ICP: \", icp_time)\n",
    "print(\"Fitness: \", reg_point_to_plane.fitness)\n",
    "print(\"Inlier RMSE: \", reg_point_to_plane.inlier_rmse)\n",
    "\n",
    "draw_registration_result(source, target, reg_point_to_plane.transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Robust ICP\n",
    "\n",
    "Using the same `max_correspondence_distance=0.5` **and** a robust kernel, we can properly register the two point clouds:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu, sigma = 0, 0.1  # mean and standard deviation\n",
    "estimation = treg.TransformationEstimationPointToPlane(\n",
    "    treg.robust_kernel.RobustKernel(\n",
    "        treg.robust_kernel.RobustKernelMethod.TukeyLoss, sigma))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Vanilla point-to-plane ICP, max_correspondence_distance={}:\".format(\n",
    "    max_correspondence_distance))\n",
    "\n",
    "s = time.time()\n",
    "\n",
    "reg_point_to_plane = treg.icp(source, target, max_correspondence_distance,\n",
    "                              init_source_to_target, estimation, criteria,\n",
    "                              voxel_size)\n",
    "\n",
    "icp_time = time.time() - s\n",
    "print(\"Time taken by Point-To-Plane ICP: \", icp_time)\n",
    "print(\"Fitness: \", reg_point_to_plane.fitness)\n",
    "print(\"Inlier RMSE: \", reg_point_to_plane.inlier_rmse)\n",
    "\n",
    "draw_registration_result(source, target, reg_point_to_plane.transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Registration estimation method supports `Robust Kernels`:\n",
    "\n",
    "`robust_kernel = o3d.t.pipelines.registration.robust_kernel.RobustKernel(method, scale, shape)`\n",
    "\n",
    "Method options:\n",
    "\n",
    "- robust_kernel.RobustKernelMethod.L2Loss\n",
    "- robust_kernel.RobustKernelMethod.L1Loss\n",
    "- robust_kernel.RobustKernelMethod.HuberLoss\n",
    "- robust_kernel.RobustKernelMethod.CauchyLoss\n",
    "- robust_kernel.RobustKernelMethod.GMLoss\n",
    "- robust_kernel.RobustKernelMethod.TukeyLoss\n",
    "- robust_kernel.RobustKernelMethod.GeneralizedLoss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "    \n",
    "**Note:** \n",
    "\n",
    "For this example we use the **TukeyLoss**. For the parameter `k` we set it to match the std deviation of the noise model $k = \\sigma$.\n",
    "The parameter `k` used in the Robust Kernels it's usually pick to match the standard deviation of the noise model of the input data. In this sense, `k`, is the discriminator between **inlier**/**outlier**. Although this is not always trivial to define in real world data, for synthetic one, it's easy to fix in order to illustrate the benefits of the robust kernels.\n",
    " \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Edit Metadata",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
