{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import open3d as o3d\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import sys\n",
    "\n",
    "# only needed for tutorial, monkey patches visualization\n",
    "sys.path.append('..')\n",
    "import open3d_tutorial\n",
    "# change to True if you want to interact with the visualization windows\n",
    "open3d_tutorial.interactive = not \"CI\" in os.environ"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distance Queries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `RaycastingScene` class in Open3D provides a set of distance queries, which can be used to convert triangle meshes into implicit functions, query the distance to the surface or determine if a point is inside a mesh."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Distance field computed from a mesh](images/distance_field_animation.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial we show how to generate these queries and produce pupular implicit representations from meshes as used in geometry processing and 3D machine learning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Converting a mesh to an implicit representation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Initialization**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As the first step we initialize a `RaycastingScene` with a (closed) triangle mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load mesh and convert to open3d.t.geometry.TriangleMesh\n",
    "armadillo_data = o3d.data.ArmadilloMesh()\n",
    "mesh = o3d.io.read_triangle_mesh(armadillo_data.path)\n",
    "mesh = o3d.t.geometry.TriangleMesh.from_legacy(mesh)\n",
    "\n",
    "# Create a scene and add the triangle mesh\n",
    "scene = o3d.t.geometry.RaycastingScene()\n",
    "_ = scene.add_triangles(mesh)  # we do not need the geometry ID for mesh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Computing distances and occupancy for a single point**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*RaycastingScene* directly provides functions to compute the unsigned and signed distance from a point to the mesh surface.\n",
    "It also provides a function to compute the occupancy at a query point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_point = o3d.core.Tensor([[10, 10, 10]], dtype=o3d.core.Dtype.Float32)\n",
    "\n",
    "# Compute distance of the query point from the surface\n",
    "unsigned_distance = scene.compute_distance(query_point)\n",
    "signed_distance = scene.compute_signed_distance(query_point)\n",
    "occupancy = scene.compute_occupancy(query_point)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While the unsigned distance can always be computed, the signed distance and the occupancy are only valid if the mesh is watertight and the inside and outside are clearly defined. \n",
    "The signed distance is negative if the query point is inside the mesh.\n",
    "The occupancy is either 0 for points outside the mesh and 1 for points inside the mesh.\n",
    "\n",
    "In this example our mesh is watertight and we can see that the query point is inside the mesh because of the signed distance and the occupancy values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"unsigned distance\", unsigned_distance.numpy())\n",
    "print(\"signed_distance\", signed_distance.numpy())\n",
    "print(\"occupancy\", occupancy.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Computing distances for multiple points and grids**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`RaycastingScene` allows multiple queries at once. \n",
    "For instance, we ca\n",
    "n pass a [N,3] Tensor with N query points which can be used to randomly sample a volume for training implicit neural representations in machine learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_bound = mesh.vertex.positions.min(0).numpy()\n",
    "max_bound = mesh.vertex.positions.max(0).numpy()\n",
    "\n",
    "N = 256\n",
    "query_points = np.random.uniform(low=min_bound, high=max_bound,\n",
    "                                 size=[N, 3]).astype(np.float32)\n",
    "\n",
    "# Compute the signed distance for N random points\n",
    "signed_distance = scene.compute_signed_distance(query_points)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`RaycastingScene` allows to organize query points with an arbitrary number of leading dimensions.\n",
    "To query the signed distance for a grid we can do the following"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xyz_range = np.linspace(min_bound, max_bound, num=32)\n",
    "\n",
    "# query_points is a [32,32,32,3] array ..\n",
    "query_points = np.stack(np.meshgrid(*xyz_range.T), axis=-1).astype(np.float32)\n",
    "\n",
    "# signed distance is a [32,32,32] array\n",
    "signed_distance = scene.compute_signed_distance(query_points)\n",
    "\n",
    "# We can visualize a slice of the distance field directly with matplotlib\n",
    "plt.imshow(signed_distance.numpy()[:, :, 15])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same procedure works for `RaycastingScene.compute_distance` and `RaycastingScene.compute_occupancy`, which can be used to generate the unsigned distance and occupancy fields."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computing distances with closest point queries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The distance functions are built on top of the `compute_closest_points()` function.\n",
    "In this part we will reimplement the signed distance and show how to make use of the additional information returned by the `compute_closest_points()` function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Initialization**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first initialize a `RaycastingScene` with two triangle meshes. \n",
    "Both meshes are watertight and we will place them such that there are no intersections between them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cube = o3d.t.geometry.TriangleMesh.from_legacy(\n",
    "    o3d.geometry.TriangleMesh.create_box().translate([-1.2, -1.2, 0]))\n",
    "sphere = o3d.t.geometry.TriangleMesh.from_legacy(\n",
    "    o3d.geometry.TriangleMesh.create_sphere(0.5).translate([0.7, 0.8, 0]))\n",
    "\n",
    "scene = o3d.t.geometry.RaycastingScene()\n",
    "# Add triangle meshes and remember ids\n",
    "mesh_ids = {}\n",
    "mesh_ids[scene.add_triangles(cube)] = 'cube'\n",
    "mesh_ids[scene.add_triangles(sphere)] = 'sphere'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Computing the closest point on the surface**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`RaycastingScene.compute_closest_points()` can compute the closest point on the surface with respect to a query point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_point = o3d.core.Tensor([[0, 0, 0]], dtype=o3d.core.Dtype.Float32)\n",
    "\n",
    "# We compute the closest point on the surface for the point at position [0,0,0].\n",
    "ans = scene.compute_closest_points(query_point)\n",
    "\n",
    "# Compute_closest_points provides the point on the surface, the geometry id,\n",
    "# and the primitive id.\n",
    "# The dictionary keys are\n",
    "#.    points\n",
    "#.    geometry_ids\n",
    "#.    primitive_ids\n",
    "print('The closest point on the surface is', ans['points'].numpy())\n",
    "print('The closest point is on the surface of the',\n",
    "      mesh_ids[ans['geometry_ids'][0].item()])\n",
    "print('The closest point belongs to triangle', ans['primitive_ids'][0].item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To compute whether the point is inside or outside we can cast a ray starting at the query point and count the number of intersections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rays = np.concatenate(\n",
    "    [query_point.numpy(),\n",
    "     np.ones(query_point.shape, dtype=np.float32)],\n",
    "    axis=-1)\n",
    "intersection_counts = scene.count_intersections(rays).numpy()\n",
    "# A point is inside if the number of intersections with the scene is even\n",
    "# This sssumes that inside and outside is we ll defined for the scene.\n",
    "is_inside = intersection_counts % 2 == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can combine this into one function to create a special signed distance function with returns additional information:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_signed_distance_and_closest_goemetry(query_points: np.ndarray):\n",
    "    closest_points = scene.compute_closest_points(query_points)\n",
    "    distance = np.linalg.norm(query_points - closest_points['points'].numpy(),\n",
    "                              axis=-1)\n",
    "    rays = np.concatenate([query_points, np.ones_like(query_points)], axis=-1)\n",
    "    intersection_counts = scene.count_intersections(rays).numpy()\n",
    "    is_inside = intersection_counts % 2 == 1\n",
    "    distance[is_inside] *= -1\n",
    "    return distance, closest_points['geometry_ids'].numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the function to create a grid with the distance and geometry id information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute range\n",
    "xyz_range = np.linspace([-2, -2, -2], [2, 2, 2], num=32)\n",
    "# query_points is a [32,32,32,3] array ..\n",
    "query_points = np.stack(np.meshgrid(*xyz_range.T), axis=-1).astype(np.float32)\n",
    "\n",
    "sdf, closest_geom = compute_signed_distance_and_closest_goemetry(query_points)\n",
    "\n",
    "# We can visualize a slice of the grids directly with matplotlib\n",
    "fig, axes = plt.subplots(1, 2)\n",
    "axes[0].imshow(sdf[:, :, 16])\n",
    "axes[1].imshow(closest_geom[:, :, 16])"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}