{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Point Cloud Registration using Standard ICP"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "In order to lower the uncertainty of calculations done with `py4dgeo`, point cloud registration should be performed to tightly align the point clouds. \n",
    "\n",
    "In this practical we will perform point cloud registration using a standard ICP implementation in `py4dgeo`. The workflow is introduced throughout this notebook. Also, make use of the [software documentation](https://py4dgeo.readthedocs.io/en/latest/pythonapi.html)!\n",
    "\n",
    "`py4dgeo` supports point cloud registration in the following way:\n",
    "\n",
    "* It allows you to apply arbitrary affine transformations to epochs and stores the transformation parameters as part of the epoch. \n",
    "* The transformation can be calculated directly in `py4dgeo` or it can be derived from an external tool and be applied to an epoch in `py4dgeo`. \n",
    "\n",
    "This notebook shows both ways of usage.\n",
    "\n",
    "`py4dgeo` returns an affine transformation as a `3x4` matrix which contains the Rotation matrix $\\mathbf{R}$ and the translation vector $\\mathbf{t}$:\n",
    "\n",
    "$$\n",
    "\\mathbf{T} = \\left(\\begin{array}{ccc}r_{11}&r_{12}&r_{13}&t_1\\\\r_{21}&r_{22}&r_{23}&t_2\\\\r_{31}&r_{32}&r_{33}&t_3\\end{array}\\right)\n",
    "$$\n",
    "\n",
    "\n",
    "Note: For the ICP to work, the two point clouds must already be roughly aligned.\n",
    "\n",
    "## Data\n",
    "We will use two TLS scans acquired at a rockfall site in Obergurgl, Austria in 2019 (epoch 1) and 2022 (epoch 2). We use only stable parts that for co-registration and later apply the estimated transformation to the full point cloud:\n",
    "- `tls_rockfall_20190619_stable_parts.laz`: Stable parts of epoch 1 (used as reference).\n",
    "- `tls_rockfall_20220921_shifted_stable_parts.laz`: Stable parts of epoch 2 (to be aligned).\n",
    "- `tls_rockfall_20220921_shifted_full_pc.laz`: Full point cloud of epoch 2 (to be aligned)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Loading data\n",
    "First, we start by setting up the Python environment and data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import py4dgeo\n",
    "import numpy as np\n",
    "import shutil\n",
    "import requests\n",
    "import tempfile\n",
    "import zipfile\n",
    "import os\n",
    "import copy\n",
    "from pathlib import Path\n",
    "\n",
    "# Url for download of data\n",
    "url = \"https://heibox.uni-heidelberg.de/f/a045ce9e0c8d4f8583d4/?dl=1\"\n",
    "\n",
    "# Path to subfolder with dataset used in this notebook\n",
    "subdir_path = \"test_data_icp\"\n",
    "\n",
    "# Generate temporary directory for data\n",
    "temp_dir = Path(tempfile.mkdtemp())\n",
    "\n",
    "try:\n",
    "    # Download data\n",
    "    response = requests.get(url)\n",
    "    response.raise_for_status()\n",
    "\n",
    "    # Store zip archive in temp dir\n",
    "    zip_path = temp_dir / \"test_data_icp.zip\"\n",
    "    with open(zip_path, \"wb\") as zip_file:\n",
    "        zip_file.write(response.content)\n",
    "\n",
    "    # Unpack zip archive in subfolder\n",
    "    with zipfile.ZipFile(zip_path, \"r\") as zip_ref:\n",
    "        file_list = [f for f in zip_ref.namelist() if f.startswith(subdir_path)]\n",
    "        zip_ref.extractall(temp_dir, members=file_list)\n",
    "\n",
    "    # specify data path\n",
    "    data_path = (temp_dir / subdir_path).resolve()\n",
    "\n",
    "except requests.exceptions.RequestException as e:\n",
    "    print(f\"Failed to download data: {e}\")\n",
    "except zipfile.BadZipFile as e:\n",
    "    print(f\"Failed to unpack data: {e}\")\n",
    "\n",
    "# check if the specified path exists\n",
    "if not os.path.isdir(data_path):\n",
    "    print(f\"ERROR: {data_path} does not exist\")\n",
    "    print(\n",
    "        \"Please specify the correct path to the data directory by replacing <path-to-data> above.\"\n",
    "    )\n",
    "\n",
    "epoch1_filepath = data_path / \"tls_rockfall_20190619_reference_stableparts.laz\"\n",
    "\n",
    "epoch2_filepath = data_path / \"tls_rockfall_20220921_shifted_stableparts.laz\"\n",
    "\n",
    "\n",
    "# Load the two epochs from LAZ files:\n",
    "epoch1 = py4dgeo.read_from_las(epoch1_filepath)\n",
    "epoch2 = py4dgeo.read_from_las(epoch1_filepath)\n",
    "\n",
    "print(f\"Epoch 1 stable parts point cloud read: {epoch1_filepath}\")\n",
    "print(f\"Epoch 2 stable parts point cloud read: {epoch2_filepath}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "Let's have a look at the documentation of the py4dgeo ICP implementation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "?py4dgeo.iterative_closest_point"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "## Definition of reduction point\n",
    "When calculating and applying transformations in `py4dgeo`, it is always possible to specify the *reduction point* $\\mathbf{x_0}$. This allows shifting coordinates towards the origin before applying rotations - leading to a numerically robust operation. The algorithm for a transformation is:\n",
    "\n",
    "$$\n",
    "\\mathbf{Tx} = \\left(\\mathbf{R}(\\mathbf{x}-\\mathbf{x_0})\\right) + \\mathbf{t} + \\mathbf{x_0}\n",
    "$$\n",
    "\n",
    "Using a reduction point is highly recommended when working with large coordinates to avoid correlation of coordinates and overflow.\n",
    "\n",
    "Since we are working with large global coordinates in this notebook, we will use a reduction point suitable for our dataset (working with your own data you might choose the coordinates of your first point in the point cloud / the center of gravity of the point cloud / etc.):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bounding box center of 2022 epoch (global coordinates)\n",
    "red_poi = np.array([653530.316040, 5191905.207520, 2020.801025])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "## Calculating transformation parameters in py4dgeo\n",
    "Now we can estimate the transformation matrix in `py4dgeo` to align both epochs. We are decreasing the tolerance and increasing the maximum number of iterations (compared to the default) to ensure very precise results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "trafo = py4dgeo.iterative_closest_point(\n",
    "    epoch1, epoch2, tolerance=0.00001, max_iterations=50\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "Let's have a look at the transformation matrix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(trafo)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "We make a copy of `epoch2` before applying the transformation to be able to use it in a later step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch2_stableparts_orig = copy.deepcopy(epoch2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "## Apply transformation\n",
    "Now, we can use this transformation matrix to transform `epoch2`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch2.transform(trafo)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "The `Epoch` class records applied transformations and makes them available through the `transformation` property:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch2.transformation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "We can save `epoch2` to a zip archive:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir = os.path.join(data_path, \"tls_rockfall_2022_registered_stable_parts.zip\")\n",
    "epoch2.save(output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "We now can apply this transformation to the full point cloud of `epoch2`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch2_full_pc = py4dgeo.read_from_las(\n",
    "    os.path.join(data_path, \"tls_rockfall_20220921_shifted.laz\")\n",
    ")\n",
    "epoch2_full_pc.transform(trafo)\n",
    "epoch2_full_pc.save(os.path.join(data_path, \"tls_rockfall_2022_registered_full_pc.zip\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "## Evaluate shift\n",
    "Let's compare the coordinates of the registered shifted point cloud in stable parts to the original:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required module\n",
    "from scipy.spatial import KDTree\n",
    "\n",
    "# Build kd-tree from 3D coordinates\n",
    "tree_orig = KDTree(epoch2_stableparts_orig.cloud)\n",
    "\n",
    "# Query indices of nearest neighbors of registered coordinates\n",
    "nn_dists = tree_orig.query(epoch2.cloud, k=1)\n",
    "\n",
    "# Obtain distances as first element in tuple returned by query above\n",
    "distances = nn_dists[0]\n",
    "\n",
    "print(f\"Mean dist: {np.mean(distances):.3f} m\")\n",
    "print(f\"Median dist: {np.median(distances):.3f} m\")\n",
    "print(f\"Std.dev. of dists: {np.std(distances):.3f} m\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "Visualize the result:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Create a figure with 3D axis (two columns for different coloring)\n",
    "fig, ax = plt.subplots(1, 1, subplot_kw={\"projection\": \"3d\"}, figsize=(7, 5))\n",
    "\n",
    "nth = 100\n",
    "\n",
    "# Plot the point cloud colored by height (z values)\n",
    "s = ax.scatter(\n",
    "    epoch2.cloud[::nth, 0],\n",
    "    epoch2.cloud[::nth, 1],\n",
    "    epoch2.cloud[::nth, 2],\n",
    "    s=5,\n",
    "    c=distances[::nth],\n",
    ")\n",
    "\n",
    "# Label axes and add title\n",
    "ax.set_xlabel(\"X [m]\")\n",
    "ax.set_ylabel(\"Y [m]\")\n",
    "ax.set_zlabel(\"Z [m]\")\n",
    "\n",
    "# Add a colorbar\n",
    "fig.colorbar(s, shrink=0.5, aspect=10, label=\"NN distance [m]\", ax=ax, pad=0.2)\n",
    "\n",
    "ax.set_aspect(\"equal\")\n",
    "\n",
    "# Show the plot\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26",
   "metadata": {},
   "source": [
    "## Applying external transformation parameters\n",
    "Having calculated a transformation matrix externally, we can apply this to epoch 2 in `py4dgeo` by loading it as numpy array. In our case, we use a transformation matrix derived in CloudCompare. \n",
    "\n",
    "Note: In CloudCompare the transformation matrix was computed based on 2 point clouds which have been [globally shifted](https://www.cloudcompare.org/doc/wiki/index.php/Global_Shift_and_Scale) at import. After applying the ICP method in CloudCompare, we saved the resulting transformation matrix without their parameters being transformed to global coordinates (second transformation matrix CloudCompare prints in the console)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "epoch2_external_trafo = py4dgeo.read_from_las(\n",
    "    os.path.join(data_path, \"tls_rockfall_20220921_shifted.laz\")\n",
    ")\n",
    "with open(os.path.join(data_path, \"external_trafo.json\"), \"r\") as json_file:\n",
    "    trafo_data = json.load(json_file)\n",
    "external_trafo = np.array(trafo_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch2_external_trafo.transform(\n",
    "    external_trafo, reduction_point=[653184.0, 5192220.0, 0]\n",
    ")\n",
    "epoch2_external_trafo.transformation\n",
    "epoch2_external_trafo.save(\n",
    "    os.path.join(data_path, \"tls_rockfall_2022_registered_external_trafo.zip\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29",
   "metadata": {},
   "source": [
    "## Available registration algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30",
   "metadata": {},
   "source": [
    "Currently only standard point to point ICP is available, but other algorithms are currently being implemented:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "?py4dgeo.iterative_closest_point"
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}