{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "import open3d as o3d\n",
    "import numpy as np\n",
    "import os\n",
    "import sys\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": [
    "# KDTree\n",
    "Open3D uses [FLANN](https://www.cs.ubc.ca/research/flann/) to build KDTrees for fast retrieval of nearest neighbors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build KDTree from point cloud\n",
    "The code below reads a point cloud and builds a KDTree. This is a preprocessing step for the following nearest neighbor queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Testing kdtree in open3d ...\")\n",
    "print(\"Load a point cloud and paint it gray.\")\n",
    "pcd = o3d.io.read_point_cloud(\"../../TestData/Feature/cloud_bin_0.pcd\")\n",
    "pcd.paint_uniform_color([0.5, 0.5, 0.5])\n",
    "pcd_tree = o3d.geometry.KDTreeFlann(pcd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Find neighboring points\n",
    "We pick the 1500-th point as the anchor point and paint it red."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Paint the 1500th point red.\")\n",
    "pcd.colors[1500] = [1, 0, 0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using search_knn_vector_3d\n",
    "The function `search_knn_vector_3d` returns a list of indices of the k nearest neighbors of the anchor point. These neighboring points are painted with blue color. Note that we convert `pcd.colors` to a numpy array to make batch access to the point colors, and broadcast a blue color [0, 0, 1] to all the selected points. We skip the first index since it is the anchor point itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Find its 200 nearest neighbors, paint blue.\")\n",
    "[k, idx, _] = pcd_tree.search_knn_vector_3d(pcd.points[1500], 200)\n",
    "np.asarray(pcd.colors)[idx[1:], :] = [0, 0, 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using search_radius_vector_3d\n",
    "Similarly, we can use `search_radius_vector_3d` to query all points with distances to the anchor point less than a given radius. We paint these points with green color."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Find its neighbors with distance less than 0.2, paint green.\")\n",
    "[k, idx, _] = pcd_tree.search_radius_vector_3d(pcd.points[1500], 0.2)\n",
    "np.asarray(pcd.colors)[idx[1:], :] = [0, 1, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Visualize the point cloud.\")\n",
    "o3d.visualization.draw_geometries([pcd], zoom=0.5599,\n",
    "                                  front=[-0.4958, 0.8229, 0.2773],\n",
    "                                  lookat=[2.1126, 1.0163, -1.8543],\n",
    "                                  up=[0.1007, -0.2626, 0.9596])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "    \n",
    "**Note:** \n",
    "\n",
    "Besides the KNN search `search_knn_vector_3d` and the RNN search `search_radius_vector_3d`, Open3D provides a hybrid search function `search_hybrid_vector_3d`. It returns at most k nearest neighbors that have distances to the anchor point less than a given radius. This function combines the criteria of KNN search and RNN search. It is known as RKNN search in some literatures. It has performance benefits in many practical cases, and is heavily used in a number of Open3D functions.\n",
    "\n",
    "</div>"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Edit 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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}