{
 "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": {
    "tags": []
   },
   "source": [
    "# ICP registration\n",
    "This tutorial demonstrates the ICP (Iterative Closest Point) registration algorithm. It has been a mainstay of geometric registration in both research and industry for many years. The inputs are two point clouds and an initial transformation that roughly aligns the source point cloud to the target point cloud. The output is a refined transformation that tightly aligns the two point clouds. A helper function `draw_registration_result` visualizes the alignment during the registration process. In this tutorial, we show different ICP variants, and the API for using them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper visualization function\n",
    "\n",
    "The function below visualizes a target point cloud and a source point cloud transformed with an alignment transformation. The target point cloud and the source point cloud are painted with cyan and yellow colors respectively. The more and tighter the two point-clouds overlap with each other, the better the alignment result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Understanding ICP Algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In general, the ICP algorithm iterates over two steps:\n",
    "\n",
    "1. Find correspondence set $\\mathcal{K}=\\{(\\mathbf{p}, \\mathbf{q})\\}$ from target point cloud $\\mathbf{P}$, and source point cloud $\\mathbf{Q}$ transformed with current transformation matrix $\\mathbf{T}$.\n",
    "2. Update the transformation $\\mathbf{T}$ by minimizing an objective function $E(\\mathbf{T})$ defined over the correspondence set $\\mathcal{K}$.\n",
    "\n",
    "Different variants of ICP use different objective functions $E(\\mathbf{T})$ [\\[BeslAndMcKay1992\\]](../reference.html#beslandmckay1992) [\\[ChenAndMedioni1992\\]](../reference.html#chenandmedioni1992) [\\[Park2017\\]](../reference.html#park2017)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understanding ICP API\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "    \n",
    "**Note:** \n",
    "\n",
    "The `Tensor based ICP implementation` API is slightly different than the [Eigen based ICP implementation](../pipelines/icp_registration.rst), to support more functionalities.\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Input\n",
    "\n",
    "\n",
    "`PointClouds` between which the `Transformation` is to be estimated. [open3d.t.PointCloud]\n",
    "\n",
    "- Source Tensor PointCloud. [Float32 or Float64 dtypes are supported].\n",
    "- Target Tensor PointCloud. [Float32 or Float64 dtypes are supported].\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "    \n",
    "**Note:** \n",
    "\n",
    "The initial alignment is usually obtained by a global registration algorithm.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "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])\n",
    "\n",
    "# For Colored-ICP `colors` attribute must be of the same dtype as `positions` and `normals` attribute.\n",
    "source.point[\"colors\"] = source.point[\"colors\"].to(\n",
    "    o3d.core.Dtype.Float32) / 255.0\n",
    "target.point[\"colors\"] = target.point[\"colors\"].to(\n",
    "    o3d.core.Dtype.Float32) / 255.0\n",
    "\n",
    "# Initial guess transform between the two point-cloud.\n",
    "# ICP algorithm requires a good initial alignment to converge efficiently.\n",
    "trans_init = 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], [0.0, 0.0, 0.0, 1.0]])\n",
    "\n",
    "draw_registration_result(source, target, trans_init)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Max correspondence Distances\n",
    "- This is the radius of distance from each point in the source point-cloud in which the neighbour search will try to find a corresponding point in the target point-cloud.\n",
    "- It is a `double` for `icp`, and `utility.DoubleVector` for `multi-scale-icp`.\n",
    "- One may typically keep this parameter between `1.0x - 3.0x` `voxel-size` for each scale.\n",
    "- This parameter is most important for performance tuning, as a higher radius will take larger time (as the neighbour search will be performed over a larger radius)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For Vanilla ICP (double)\n",
    "\n",
    "# Search distance for Nearest Neighbour Search [Hybrid-Search is used].\n",
    "max_correspondence_distance = 0.07"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For Multi-Scale ICP (o3d.utility.DoubleVector):\n",
    "\n",
    "# `max_correspondence_distances` is proportional to the resolution or the `voxel_sizes`.\n",
    "# In general it is recommended to use values between 1x - 3x of the corresponding `voxel_sizes`.\n",
    "# We may have a higher value of the `max_correspondence_distances` for the first coarse\n",
    "# scale, as it is not much expensive, and gives us more tolerance to initial alignment.\n",
    "max_correspondence_distances = o3d.utility.DoubleVector([0.3, 0.14, 0.07])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initial Transform from Source to Target [open3d.core.Tensor]\n",
    "- Initial estimate for transformation from source to target.\n",
    "- Transformation matrix Tensor of shape [4, 4] of type `Float64` on `CPU:0` device\n",
    "- The initial alignment is usually obtained by a global registration algorithm. See [Global registration](../pipelines/global_registration.rst) for examples."
   ]
  },
  {
   "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]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Estimation Method \n",
    "- This sets the ICP method to compute the transformation between two point-clouds given the correspondences.\n",
    "\n",
    "Options:\n",
    "\n",
    "- **o3d.t.pipelines.registration.TransformationEstimationPointToPoint()**\n",
    "    - Point to Point ICP.\n",
    "- **o3d.t.pipelines.registration.TransformationEstimationPointToPlane(robust_kernel)**\n",
    "    - Point to Plane ICP.\n",
    "    - Requires `target point-cloud` to have `normals` attribute (of same dtype as `position` attribute).\n",
    "- **o3d.t.pipelines.registration.TransformationEstimationForColoredICP(robust_kernel, lambda)**\n",
    "    - Colored ICP.\n",
    "    - Requires `target` point-cloud to have `normals` attribute (of same dtype as `position` attribute).\n",
    "    - Requires `source` and `target` point-clouds to have `colors` attribute (of same dtype as `position` attribute).\n",
    "- **o3d.t.pipelines.registration.TransformationEstimationForGeneralizedICP(robust_kernel, epsilon)** [To be added].\n",
    "    - Generalized ICP.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select the `Estimation Method`, and `Robust Kernel` (for outlier-rejection).\n",
    "estimation = treg.TransformationEstimationPointToPlane()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Estimation Method also supports `Robust Kernels`: Robust kernels are used for outlier rejection. More on this in `Robust Kernel` section.\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\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### ICP Convergence Criteria [relative rmse, relative fitness, max iterations]\n",
    "- This sets the condition for termination or when the scale iterations can be considered to be converged. \n",
    "- If the relative (of change in value from the last iteration) rmse and fitness are equal or less than the specified value, the iterations for that scale will be considered as converged/completed.\n",
    "- For `Multi-Scale ICP` it is a `list` of `ICPConvergenceCriteria`, for each scale of ICP, to provide more fine control over performance.\n",
    "- One may keep the initial values of `relative_fitness` and `relative_rmse` low as we just want to get an estimate transformation, and high for later iterations to fine-tune.\n",
    "- Iterations on higher-resolution is more costly (takes more time), so we want to do fewer iterations on higher resolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convergence-Criteria for Vanilla ICP:\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": [
    "# List of Convergence-Criteria for Multi-Scale ICP:\n",
    "\n",
    "# We can control `ConvergenceCriteria` of each `scale` individually.\n",
    "# We want to keep `relative_fitness` and `relative_rmse` high (more error tolerance)\n",
    "# for initial scales, i.e. we will be happy to consider ICP converged, when difference\n",
    "# between 2 successive iterations for that scale is smaller than this value.\n",
    "# We expect less accuracy (more error tolerance) initial coarse-scale iteration,\n",
    "# and want our later scale convergence to be more accurate (less error tolerance).\n",
    "criteria_list = [\n",
    "    treg.ICPConvergenceCriteria(relative_fitness=0.0001,\n",
    "                                relative_rmse=0.0001,\n",
    "                                max_iteration=20),\n",
    "    treg.ICPConvergenceCriteria(0.00001, 0.00001, 15),\n",
    "    treg.ICPConvergenceCriteria(0.000001, 0.000001, 10)\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Voxel Sizes\n",
    "- It is the voxel size (lower voxel size corresponds to higher resolution), for each scale of multi-scale ICP.\n",
    "- We want to perform initial iterations on a coarse point-cloud (low-resolution or small voxel size) (as it is more time-efficient, and avoids local-minima), and then move to a dense point-cloud (high-resolution or small voxel size. Therefore the voxel sizes must be strictly decreasing order."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Vanilla ICP\n",
    "voxel_size = 0.025"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lower `voxel_size` is equivalent to higher resolution,\n",
    "# and we want to perform iterations from coarse to dense resolution,\n",
    "# therefore `voxel_sizes` must be in strictly decressing order.\n",
    "voxel_sizes = o3d.utility.DoubleVector([0.1, 0.05, 0.025])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Get Iteration-wise registration result using callback lambda function\n",
    "\n",
    "Optional lambda function, saves string to tensor dictionary of attributes such as \"iteration_index\", \"scale_index\", \"scale_iteration_index\", \"inlier_rmse\", \"fitness\", \"transformation\", on CPU device, updated after each iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example callback_after_iteration lambda function:\n",
    "callback_after_iteration = lambda updated_result_dict : print(\"Iteration Index: {}, Fitness: {}, Inlier RMSE: {},\".format(\n",
    "    updated_result_dict[\"iteration_index\"].item(),\n",
    "    updated_result_dict[\"fitness\"].item(), \n",
    "    updated_result_dict[\"inlier_rmse\"].item()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vanilla ICP Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Set Inputs and Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input point-clouds\n",
    "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": [
    "# Search distance for Nearest Neighbour Search [Hybrid-Search is used].\n",
    "max_correspondence_distance = 0.07\n",
    "\n",
    "# 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",
    "# Select the `Estimation Method`, and `Robust Kernel` (for outlier-rejection).\n",
    "estimation = treg.TransformationEstimationPointToPlane()\n",
    "\n",
    "# Convergence-Criteria for Vanilla ICP\n",
    "criteria = treg.ICPConvergenceCriteria(relative_fitness=0.000001,\n",
    "                                       relative_rmse=0.000001,\n",
    "                                       max_iteration=50)\n",
    "# Down-sampling voxel-size.\n",
    "voxel_size = 0.025\n",
    "\n",
    "# Save iteration wise `fitness`, `inlier_rmse`, etc. to analyse and tune result.\n",
    "save_loss_log = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Get Registration Result from ICP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s = time.time()\n",
    "\n",
    "registration_icp = treg.icp(source, target, max_correspondence_distance,\n",
    "                            init_source_to_target, estimation, criteria,\n",
    "                            voxel_size, callback_after_iteration)\n",
    "\n",
    "icp_time = time.time() - s\n",
    "print(\"Time taken by ICP: \", icp_time)\n",
    "print(\"Inlier Fitness: \", registration_icp.fitness)\n",
    "print(\"Inlier RMSE: \", registration_icp.inlier_rmse)\n",
    "\n",
    "draw_registration_result(source, target, registration_icp.transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "---\n",
    "Now let's try with poor initial initialisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_source_to_target = o3d.core.Tensor.eye(4, o3d.core.Dtype.Float32)\n",
    "max_correspondence_distance = 0.07"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s = time.time()\n",
    "\n",
    "registration_icp = 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 ICP: \", icp_time)\n",
    "print(\"Inlier Fitness: \", registration_icp.fitness)\n",
    "print(\"Inlier RMSE: \", registration_icp.inlier_rmse)\n",
    "\n",
    "draw_registration_result(source, target, registration_icp.transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, poor initial alignment might fail ICP convergence\n",
    "\n",
    "Having large `max_correspondence_distance` might resolve this issue. But it will take longer to process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_source_to_target = o3d.core.Tensor.eye(4, o3c.float32)\n",
    "max_correspondence_distance = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s = time.time()\n",
    "# It is highly recommended to down-sample the point-cloud before using\n",
    "# ICP algorithm, for better performance.\n",
    "\n",
    "registration_icp = 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 ICP: \", icp_time)\n",
    "print(\"Inlier Fitness: \", registration_icp.fitness)\n",
    "print(\"Inlier RMSE: \", registration_icp.inlier_rmse)\n",
    "\n",
    "draw_registration_result(source, target, registration_icp.transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We may resolve the above issues and get even better accuracy by using `Multi-Scale ICP`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Multi-Scale ICP Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Problems with using Vanilla-ICP (previous version):\n",
    "\n",
    "- Running the ICP algorithm on dense point-clouds is very slow. \n",
    "- It requires good initial alignment:\n",
    "    - If the point-cloud is not well aligned, the convergence might get stuck in local-minima in initial iterations.\n",
    "    - We need to have a larger `max_correspondence_distance` if the aligned point cloud does not have sufficient overlaps.  \n",
    "- If point-cloud is heavily down-sampled (coarse), the obtained result will not be accurate. \n",
    "\n",
    "These drawbacks can be solved using Multi-Scale ICP.\n",
    "In Multi-Scale ICP, we perform the initial iterations on coarse point-cloud to get a better estimate of initial alignment and use this alignment for convergence on a more dense point cloud. ICP on coarse point cloud is in-expensive, and allows us to use a larger `max_correspondence_distance`. It is also less likely for the convergence to get stuck in local minima. As we get a good estimate, it takes fewer iterations on dense point-cloud to converge to a more accurate transform. \n",
    "\n",
    "It is recommended to use `Multi-Scale ICP` over `ICP`, for efficient convergence, especially for large point clouds. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Set Inputs and Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input point-clouds\n",
    "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": [
    "voxel_sizes = o3d.utility.DoubleVector([0.1, 0.05, 0.025])\n",
    "\n",
    "# List of Convergence-Criteria for Multi-Scale ICP:\n",
    "criteria_list = [\n",
    "    treg.ICPConvergenceCriteria(relative_fitness=0.0001,\n",
    "                                relative_rmse=0.0001,\n",
    "                                max_iteration=20),\n",
    "    treg.ICPConvergenceCriteria(0.00001, 0.00001, 15),\n",
    "    treg.ICPConvergenceCriteria(0.000001, 0.000001, 10)\n",
    "]\n",
    "\n",
    "# `max_correspondence_distances` for Multi-Scale ICP (o3d.utility.DoubleVector):\n",
    "max_correspondence_distances = o3d.utility.DoubleVector([0.3, 0.14, 0.07])\n",
    "\n",
    "# Initial alignment or source to target transform.\n",
    "init_source_to_target = o3d.core.Tensor.eye(4, o3d.core.Dtype.Float32)\n",
    "\n",
    "# Select the `Estimation Method`, and `Robust Kernel` (for outlier-rejection).\n",
    "estimation = treg.TransformationEstimationPointToPlane()\n",
    "\n",
    "# Save iteration wise `fitness`, `inlier_rmse`, etc. to analyse and tune result.\n",
    "callback_after_iteration = lambda loss_log_map : print(\"Iteration Index: {}, Scale Index: {}, Scale Iteration Index: {}, Fitness: {}, Inlier RMSE: {},\".format(\n",
    "    loss_log_map[\"iteration_index\"].item(), \n",
    "    loss_log_map[\"scale_index\"].item(), \n",
    "    loss_log_map[\"scale_iteration_index\"].item(), \n",
    "    loss_log_map[\"fitness\"].item(), \n",
    "    loss_log_map[\"inlier_rmse\"].item()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Get Registration Result from Multi-Scale ICP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting Verbosity to Debug, helps in fine-tuning the performance.\n",
    "# o3d.utility.set_verbosity_level(o3d.utility.VerbosityLevel.Debug)\n",
    "\n",
    "s = time.time()\n",
    "\n",
    "registration_ms_icp = treg.multi_scale_icp(source, target, voxel_sizes,\n",
    "                                           criteria_list,\n",
    "                                           max_correspondence_distances,\n",
    "                                           init_source_to_target, estimation,\n",
    "                                           callback_after_iteration)\n",
    "\n",
    "ms_icp_time = time.time() - s\n",
    "print(\"Time taken by Multi-Scale ICP: \", ms_icp_time)\n",
    "print(\"Inlier Fitness: \", registration_ms_icp.fitness)\n",
    "print(\"Inlier RMSE: \", registration_ms_icp.inlier_rmse)\n",
    "\n",
    "draw_registration_result(source, target, registration_ms_icp.transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-Scale ICP on CUDA device Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The algorithm runs on the same device as the source and target point-cloud.\n",
    "source_cuda = source.cuda(0)\n",
    "target_cuda = target.cuda(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s = time.time()\n",
    "\n",
    "registration_ms_icp = treg.multi_scale_icp(source_cuda, target_cuda,\n",
    "                                           voxel_sizes, criteria_list,\n",
    "                                           max_correspondence_distances,\n",
    "                                           init_source_to_target, estimation)\n",
    "\n",
    "ms_icp_time = time.time() - s\n",
    "print(\"Time taken by Multi-Scale ICP: \", ms_icp_time)\n",
    "print(\"Inlier Fitness: \", registration_ms_icp.fitness)\n",
    "print(\"Inlier RMSE: \", registration_ms_icp.inlier_rmse)\n",
    "\n",
    "draw_registration_result(source.cpu(), target.cpu(),\n",
    "                         registration_ms_icp.transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Case of `no correspondences`.\n",
    "In case of `no_correspondences` the `fitness` and `inlier_rmse` is `0`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_correspondence_distance = 0.02\n",
    "\n",
    "init_source_to_target = np.asarray([[1.0, 0.0, 0.0, 5], [0.0, 1.0, 0.0, 7],\n",
    "                                    [0.0, 0.0, 1.0, 10], [0.0, 0.0, 0.0, 1.0]])\n",
    "\n",
    "registration_icp = treg.icp(source, target, max_correspondence_distance,\n",
    "                            init_source_to_target)\n",
    "\n",
    "print(\"Inlier Fitness: \", registration_icp.fitness)\n",
    "print(\"Inlier RMSE: \", registration_icp.inlier_rmse)\n",
    "print(\"Transformation: \\n\", registration_icp.transformation)\n",
    "\n",
    "if registration_icp.fitness == 0 and registration_icp.inlier_rmse == 0:\n",
    "    print(\"ICP Convergence Failed, as no correspondence were found\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Information Matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Information Matrix` gives us further information about how well the point-clouds are aligned. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "information_matrix = treg.get_information_matrix(\n",
    "    source, target, max_correspondence_distances[2],\n",
    "    registration_ms_icp.transformation)\n",
    "\n",
    "print(information_matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Now that we have a basic understanding of the ICP algorithm and the API, let's experiment with the different versions to understand the difference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initial Alignment"
   ]
  },
  {
   "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])\n",
    "\n",
    "# Initial guess transform between the two point-cloud.\n",
    "# ICP algorithm requires a good initial alignment to converge efficiently.\n",
    "trans_init = 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], [0.0, 0.0, 0.0, 1.0]])\n",
    "\n",
    "draw_registration_result(source, target, trans_init)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Search distance for Nearest Neighbour Search [Hybrid-Search is used].\n",
    "max_correspondence_distance = 0.02\n",
    "\n",
    "print(\"Initial alignment\")\n",
    "evaluation = treg.evaluate_registration(source, target,\n",
    "                                        max_correspondence_distance, trans_init)\n",
    "\n",
    "print(\"Fitness: \", evaluation.fitness)\n",
    "print(\"Inlier RMSE: \", evaluation.inlier_rmse)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Point-To-Point ICP Registration \n",
    "\n",
    "We first show a point-to-point ICP algorithm [\\[BeslAndMcKay1992\\]](../reference.html#beslandmckay1992) using the objective\n",
    "\n",
    "\\begin{equation}\n",
    "E(\\mathbf{T}) = \\sum_{(\\mathbf{p},\\mathbf{q})\\in\\mathcal{K}}\\|\\mathbf{p} - \\mathbf{T}\\mathbf{q}\\|^{2}\n",
    "\\end{equation}\n",
    "\n",
    "The class `TransformationEstimationPointToPoint` provides functions to compute the residuals and Jacobian matrices of the point-to-point ICP objective."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Set Inputs and Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input point-clouds\n",
    "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": [
    "# Select the `Estimation Method`, and `Robust Kernel` (for outlier-rejection).\n",
    "estimation = treg.TransformationEstimationPointToPoint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Search distance for Nearest Neighbour Search [Hybrid-Search is used].\n",
    "max_correspondence_distance = 0.02\n",
    "\n",
    "# Initial alignment or source to target transform.\n",
    "init_source_to_target = trans_init\n",
    "\n",
    "# Convergence-Criteria for Vanilla ICP\n",
    "criteria = treg.ICPConvergenceCriteria(relative_fitness=0.0000001,\n",
    "                                       relative_rmse=0.0000001,\n",
    "                                       max_iteration=30)\n",
    "\n",
    "# Down-sampling voxel-size. If voxel_size < 0, original scale is used.\n",
    "voxel_size = -1\n",
    "\n",
    "# Save iteration wise `fitness`, `inlier_rmse`, etc. to analyse and tune result.\n",
    "save_loss_log = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Apply Point-to-Point ICP\")\n",
    "s = time.time()\n",
    "\n",
    "reg_point_to_point = 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-Point ICP: \", icp_time)\n",
    "print(\"Fitness: \", reg_point_to_point.fitness)\n",
    "print(\"Inlier RMSE: \", reg_point_to_point.inlier_rmse)\n",
    "\n",
    "draw_registration_result(source, target, reg_point_to_point.transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fitness score increases from `0.174722` to `0.372474`. The inlier_rmse reduces from `0.011771` to `0.007761`. \n",
    "By default, icp runs until convergence or reaches a maximum number of iterations (30 by default). \n",
    "It can be changed to allow more computation time and to improve the results further."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "criteria = treg.ICPConvergenceCriteria(relative_fitness=0.0000001,\n",
    "                                       relative_rmse=0.0000001,\n",
    "                                       max_iteration=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Apply Point-to-Point ICP\")\n",
    "s = time.time()\n",
    "\n",
    "reg_point_to_point = 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-Point ICP: \", icp_time)\n",
    "print(\"Fitness: \", reg_point_to_point.fitness)\n",
    "print(\"Inlier RMSE: \", reg_point_to_point.inlier_rmse)\n",
    "\n",
    "draw_registration_result(source, target, reg_point_to_point.transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The final alignment is tight. The fitness score improves to `0.620972`. The inlier_rmse reduces to `0.006581`.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Point-to-Plane ICP Registration\n",
    "The point-to-plane ICP algorithm [\\[ChenAndMedioni1992\\]](../reference.html#chenandmedioni1992) uses a different 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}$. [\\[Rusinkiewicz2001\\]](../reference.html#rusinkiewicz2001) has shown that the point-to-plane ICP algorithm has a faster convergence speed than the point-to-point ICP algorithm.\n",
    "\n",
    "The class `TransformationEstimationPointToPlane` provides functions to compute the residuals and Jacobian matrices of the point-to-plane ICP objective. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimation = treg.TransformationEstimationPointToPlane()\n",
    "\n",
    "criteria = treg.ICPConvergenceCriteria(relative_fitness=0.0000001,\n",
    "                                       relative_rmse=0.0000001,\n",
    "                                       max_iteration=30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Apply Point-to-Plane ICP\")\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": [
    "The point-to-plane ICP reaches tight alignment within 30 iterations (a `fitness` score of 0.620972 and an `inlier_rmse` score of 0.006581)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Colored ICP Registration\n",
    "This tutorial demonstrates an ICP variant that uses both geometry and color for registration. It implements the algorithm of [\\[Park2017\\]](../reference.html#park2017). The color information locks the alignment along the tangent plane. Thus this algorithm is more accurate and more robust than prior point cloud registration algorithms, while the running speed is comparable to that of ICP registration. "
   ]
  },
  {
   "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.5,\n",
    "        front=[-0.2458, -0.8088, 0.5342],\n",
    "        lookat=[1.7745, 2.2305, 0.9787],\n",
    "        up=[0.3109, -0.5878, -0.7468])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"1. Load two point clouds and show initial pose\")\n",
    "demo_cicp_pcds = o3d.data.DemoColoredICPPointClouds()\n",
    "source = o3d.t.io.read_point_cloud(demo_cicp_pcds.paths[0])\n",
    "target = o3d.t.io.read_point_cloud(demo_cicp_pcds.paths[1])\n",
    "\n",
    "# For Colored-ICP `colors` attribute must be of the same dtype as `positions` and `normals` attribute.\n",
    "source.point[\"colors\"] = source.point[\"colors\"].to(\n",
    "    o3d.core.Dtype.Float32) / 255.0\n",
    "target.point[\"colors\"] = target.point[\"colors\"].to(\n",
    "    o3d.core.Dtype.Float32) / 255.0\n",
    "\n",
    "# draw initial alignment\n",
    "current_transformation = np.identity(4)\n",
    "draw_registration_result(source, target, current_transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting baseline with point-to-plane registration\n",
    "We first run Point-to-plane ICP as a baseline approach. The visualization below shows misaligned green triangle textures. This is because a geometric constraint does not prevent two planar surfaces from slipping."
   ]
  },
  {
   "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\n",
    "init_source_to_target = np.identity(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Apply Point-to-Plane ICP\")\n",
    "s = time.time()\n",
    "\n",
    "reg_point_to_plane = treg.icp(source, target, max_correspondence_distance,\n",
    "                              init_source_to_target, estimation)\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": {
    "tags": []
   },
   "source": [
    "## Colored Registration\n",
    "The core function for colored point cloud registration is `registration_colored_icp`. Following [\\[Park2017\\]](../reference.html#park2017), it runs ICP iterations (see [Point-to-point ICP](../pipelines/icp_registration.ipynb#Point-to-point-ICP) for details) with a joint optimization objective\n",
    "\n",
    "\\begin{equation}\n",
    "E(\\mathbf{T}) = (1-\\delta)E_{C}(\\mathbf{T}) + \\delta E_{G}(\\mathbf{T})\n",
    "\\end{equation}\n",
    "\n",
    "where $\\mathbf{T}$ is the transformation matrix to be estimated. $E_{C}$ and $E_{G}$ are the photometric and geometric terms, respectively. $\\delta\\in[0,1]$ is a weight parameter that has been determined empirically.\n",
    "\n",
    "The geometric term $E_{G}$ is the same as the [Point-to-plane ICP](../pipelines/icp_registration.ipynb#Point-to-plane-ICP) objective\n",
    "\n",
    "\\begin{equation}\n",
    "E_{G}(\\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 $\\mathcal{K}$ is the correspondence set in the current iteration. $\\mathbf{n}_{\\mathbf{p}}$ is the normal of point $\\mathbf{p}$.\n",
    "\n",
    "The color term $E_{C}$ measures the difference between the color of point $\\mathbf{q}$ (denoted as $C(\\mathbf{q})$) and the color of its projection on the tangent plane of $\\mathbf{p}$.\n",
    "\n",
    "\\begin{equation}\n",
    "E_{C}(\\mathbf{T}) = \\sum_{(\\mathbf{p},\\mathbf{q})\\in\\mathcal{K}}\\big(C_{\\mathbf{p}}(\\mathbf{f}(\\mathbf{T}\\mathbf{q})) - C(\\mathbf{q})\\big)^{2},\n",
    "\\end{equation}\n",
    "\n",
    "where $C_{\\mathbf{p}}(\\cdot)$ is a precomputed function continuously defined on the tangent plane of $\\mathbf{p}$. Function$\\mathbf{f}(\\cdot)$ projects a 3D point to the tangent plane. For more details, refer to [\\[Park2017\\]](../reference.html#park2017).\n",
    "\n",
    "To further improve efficiency, [\\[Park2017\\]](../reference.html#park2017) proposes a multi-scale registration scheme. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimation = treg.TransformationEstimationForColoredICP()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_transformation = np.identity(4)\n",
    "\n",
    "criteria_list = [\n",
    "    treg.ICPConvergenceCriteria(relative_fitness=0.0001,\n",
    "                                relative_rmse=0.0001,\n",
    "                                max_iteration=50),\n",
    "    treg.ICPConvergenceCriteria(0.00001, 0.00001, 30),\n",
    "    treg.ICPConvergenceCriteria(0.000001, 0.000001, 14)\n",
    "]\n",
    "\n",
    "max_correspondence_distances = o3d.utility.DoubleVector([0.08, 0.04, 0.02])\n",
    "\n",
    "voxel_sizes = o3d.utility.DoubleVector([0.04, 0.02, 0.01])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# colored pointcloud registration\n",
    "# This is implementation of following paper\n",
    "# J. Park, Q.-Y. Zhou, V. Koltun,\n",
    "# Colored Point Cloud Registration Revisited, ICCV 2017\n",
    "\n",
    "print(\"Colored point cloud registration\")\n",
    "s = time.time()\n",
    "\n",
    "reg_multiscale_icp = treg.multi_scale_icp(source, target, voxel_sizes,\n",
    "                                          criteria_list,\n",
    "                                          max_correspondence_distances,\n",
    "                                          init_source_to_target, estimation)\n",
    "\n",
    "icp_time = time.time() - s\n",
    "print(\"Time taken by Colored 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_multiscale_icp.transformation)"
   ]
  }
 ],
 "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
}
