{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "import open3d as o3d\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import copy\n",
    "import os\n",
    "import sys\n",
    "\n",
    "# only needed for tutorial, monkey patches visualization\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": [
    "# Octree\n",
    "An **octree** is a tree data structure where each internal node has eight children. Octrees are commonly used for spatial partitioning of 3D point clouds. Non-empty leaf nodes of an octree contain one or more points that fall within the same spatial subdivision. Octrees are a useful description of 3D space and can be used to quickly find nearby points. Open3D has the geometry type `Octree` that can be used to create, search, and traverse octrees with a user-specified maximum tree depth, `max_depth`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## From point cloud\n",
    "An octree can be constructed from a point cloud using the method `convert_from_point_cloud`. Each point is inserted into the tree by following the path from the root node to the appropriate leaf node at depth `max_depth`. As the tree depth increases, internal (and eventually leaf) nodes represents a smaller partition of 3D space.\n",
    "\n",
    "If the point cloud has color, the the corresponding leaf node takes the color of the last inserted point. The `size_expand` parameter increases the size of the root octree node so it is slightly bigger than the original point cloud bounds to accommodate all points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('input')\n",
    "N = 2000\n",
    "\n",
    "armadillo = o3d.data.ArmadilloMesh()\n",
    "mesh = o3d.io.read_triangle_mesh(armadillo.path)\n",
    "pcd = mesh.sample_points_poisson_disk(N)\n",
    "# fit to unit cube\n",
    "pcd.scale(1 / np.max(pcd.get_max_bound() - pcd.get_min_bound()),\n",
    "          center=pcd.get_center())\n",
    "pcd.colors = o3d.utility.Vector3dVector(np.random.uniform(0, 1, size=(N, 3)))\n",
    "o3d.visualization.draw_geometries([pcd])\n",
    "\n",
    "print('octree division')\n",
    "octree = o3d.geometry.Octree(max_depth=4)\n",
    "octree.convert_from_point_cloud(pcd, size_expand=0.01)\n",
    "o3d.visualization.draw_geometries([octree])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## From voxel grid\n",
    "An octree can also be constructed from an Open3D `VoxelGrid` geometry using the method `create_from_voxel_grid`. Each voxel of the input `VoxelGrid` is treated as a point in 3D space with coordinates corresponding to the origin of the voxel. Each leaf node takes the color of its corresponding voxel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('voxelization')\n",
    "voxel_grid = o3d.geometry.VoxelGrid.create_from_point_cloud(pcd,\n",
    "                                                            voxel_size=0.05)\n",
    "o3d.visualization.draw_geometries([voxel_grid])\n",
    "\n",
    "print('octree division')\n",
    "octree = o3d.geometry.Octree(max_depth=4)\n",
    "octree.create_from_voxel_grid(voxel_grid)\n",
    "o3d.visualization.draw_geometries([octree])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally, an `Octree` can be converted to a `VoxelGrid` with `to_voxel_grid`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Traversal\n",
    "An octree can be traversed which can be useful for searching or processing subsections of 3D geometry. By providing the `traverse` method with a callback, each time a node (internal or leaf) is visited, additional processing can be performed.\n",
    "\n",
    "In the following example, an early stopping criterion is used to only process internal/leaf nodes with more than a certain number of points. This early stopping ability can be used to efficiently process spatial regions meeting certain conditions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_traverse(node, node_info):\n",
    "    early_stop = False\n",
    "\n",
    "    if isinstance(node, o3d.geometry.OctreeInternalNode):\n",
    "        if isinstance(node, o3d.geometry.OctreeInternalPointNode):\n",
    "            n = 0\n",
    "            for child in node.children:\n",
    "                if child is not None:\n",
    "                    n += 1\n",
    "            print(\n",
    "                \"{}{}: Internal node at depth {} has {} children and {} points ({})\"\n",
    "                .format('    ' * node_info.depth,\n",
    "                        node_info.child_index, node_info.depth, n,\n",
    "                        len(node.indices), node_info.origin))\n",
    "\n",
    "            # we only want to process nodes / spatial regions with enough points\n",
    "            early_stop = len(node.indices) < 250\n",
    "    elif isinstance(node, o3d.geometry.OctreeLeafNode):\n",
    "        if isinstance(node, o3d.geometry.OctreePointColorLeafNode):\n",
    "            print(\"{}{}: Leaf node at depth {} has {} points with origin {}\".\n",
    "                  format('    ' * node_info.depth, node_info.child_index,\n",
    "                         node_info.depth, len(node.indices), node_info.origin))\n",
    "    else:\n",
    "        raise NotImplementedError('Node type not recognized!')\n",
    "\n",
    "    # early stopping: if True, traversal of children of the current node will be skipped\n",
    "    return early_stop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "octree = o3d.geometry.Octree(max_depth=4)\n",
    "octree.convert_from_point_cloud(pcd, size_expand=0.01)\n",
    "octree.traverse(f_traverse)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Find leaf node containing point\n",
    "Using the above traversal mechanism, an octree can be quickly searched for the leaf node that contains a given point. This functionality is provided via the `locate_leaf_node` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "octree.locate_leaf_node(pcd.points[0])"
   ]
  }
 ],
 "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}