RC3DR_no_out.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RC3DR\n",
    "Robot-controlled 3D Reconstruction\n",
    "\n",
    "\n",
    "This project is heavily based on the following two papers:\n",
    "- Xinyi Fan, Linguang Zhang, Benedict Brown, and Szymon Rusinkiewicz. 2016. Automated view and path planning for scalable multi-object 3D scanning. ACM Trans. Graph. 35, 6, Article 239 (November 2016), 13 pages. DOI:https://doi.org/10.1145/2980179.2980225\n",
    "- R. A. Newcombe et al., \"KinectFusion: Real-time dense surface mapping and tracking,\" 2011 10th IEEE International Symposium on Mixed and Augmented Reality, Basel, 2011, pp. 127-136.\n",
    "\n",
    "\n",
    "## *Required packages*:\n",
    "- numpy\n",
    "- scipy\n",
    "- numba\n",
    "- python-socketio\n",
    "- tornado\n",
    "- skimage\n",
    "- igl\n",
    "- pyquaternion\n",
    "\n",
    "![blubb](res/Poster.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "This project uses modern web technologies to interact with the program instead of traditional GUI frameworks. In the future, it should be possible to work from any mobile device at any location and share results with the world quickly and without any setup.\n",
    "\n",
    "First, we have to start the server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "\n",
    "# Server\n",
    "import asyncio # event loop\n",
    "from dash import Dash # Web-App\n",
    "\n",
    "# rc3dr\n",
    "from rc3dr import rc3dr, utils, view_traversal\n",
    "\n",
    "# processing\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "import igl\n",
    "\n",
    "# cuda kernel\n",
    "from numba import cuda\n",
    "import math\n",
    "\n",
    "# plotting\n",
    "import ipyvolume as ipv\n",
    "import matplotlib.pyplot as plt \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# specify port on where to run the server\n",
    "port = 8005\n",
    "# initialize Dash\n",
    "dash = Dash(port) \n",
    "\n",
    "# Since this notebook runs on an event loop we have to create a seperate task for the server\n",
    "loop = asyncio.get_event_loop()\n",
    "async def start_server(dash):\n",
    "    dash.start_server()\n",
    "loop.create_task(start_server(dash)) # Coroutine gets killed as soon as kernel of this notebook stops"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, go to *localhost:port*. You should see something like this:\n",
    "\n",
    "<img src=\"./res/dash-image.jpg\" height=60>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's start\n",
    "\n",
    "On the *Dash* we can take various *actions* that lead us through the 3D reconstruction pipeline:\n",
    "\n",
    "- **Start TSDF**: Start capturing depth images and send those to the server. On the server we compute a *truncated signed distance function (TSDF)* from the depth image and *fuse* it together with the TSDFs obtained from previous depth images.\n",
    "- **Start exploration**: The 3D reconstruction pipeline starts with a scene exploration step to capture an initial 3D estimate of the target object. In the subsequent steps, we use this information to compute views that heuristically contain the most information about the object.\n",
    "- **Sample views**: We compute those views by first sampling views on a sphere around the object in order to keep a save distance from the camera.\n",
    "- **Select views**: Next, we select a minimal subset of those views that together contain the most information about the object. This step has multiple substeps: First, we reconstruct a mesh from the TSDF (marching cubes) and process the resulting mesh (smoothing). Second, we compute the curvature for each vertex in the mesh and sample vertices weighted by their curvature. This gives more weight to surface areas that contain more details. In the third step, we compute for each view and each sampled vertex a *geometry score* (angle between ray from vertex to camera and surface normal) and a *visibility value* (1 if in view space, 0 otherwise). Finally, we successively select views that improve the objective function $f(min\\ subset)$ most at each iteration (we terminate if there is no further improvement). \n",
    "$$ f(\\text{min subset}) = \\sum_{p\\ \\in\\ \\text{sampled vertices}}{max_{v\\ \\in\\ \\text{min subset}}{\\text{geometry}(p,v) \\cdot \\text{visibility}(p,v)}} - \\gamma \\cdot \\vert \\text{min subset} \\vert$$\n",
    "Note, that the set-cover problem is NP-complete and in most cases requires approximating algorithms.\n",
    "- **Traverse views**: We are left with finding the shortest path that connects our chosen views. This corresponds to the travelling salesman problem that is NP-hard. Here, we use a greedy approach and select at each step the view that is nearest to the current one. We define the distance $d(a,b)$ between two views $a$ and $b$ as $ d(a,b)= \\max{\\vert q_a - q_b \\vert} $, where $q_a$ and $q_b$ are the robot generalized coordinates for views $a$ and $b$ (in our case the angles of the revolute joints).\n",
    "- **Flip volume**: Rotate the TSDF volume and target object by 180 degree around the y-axis. This very naively emulates a second robot arm that can pick the target object from the other side.\n",
    "\n",
    "Now it's time to experiment. Press **Start TSDF** and then **Start exploration**. After the robot completeted its movement press **Stop TSDF** and execute the subsequent cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "F = dash.rc3dr.get_tsdf_volume()\n",
    "verts, faces = dash.rc3dr.view_selector.tsdf_to_mesh(F)\n",
    "\n",
    "ipv.figure()\n",
    "mesh = ipv.plot_trisurf(verts[:,0], verts[:,1], verts[:,2], triangles=faces, color='violet')\n",
    "ipv.squarelim()\n",
    "ipv.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also inspect the TSDF directly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Try to update slider of first level.\n",
    "# NOTE Widget needs some time to load.\n",
    "ipv.quickvolshow(F, level=0, opacity=1, level_width=0.1, data_min=-1, data_max=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now go to the *Dash* and click you through the 3D reconstruction pipeline detailed above. Don't forget to start/stop the depth image recording. Then execute the cells again."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code Details\n",
    "\n",
    "The RC3DR code is structured into following files:\n",
    "- rc3dr.py: Store data (e.g. TSDF, robot etc.) and handle IO.\n",
    "- fusion.py: TSDF fusion.\n",
    "- object_exploration.py: Generate object exploration steps.\n",
    "- view_sampling.py: Sample views.\n",
    "- view_selection.py: Select views.\n",
    "- view_traversal.py: Generate view traversal steps.\n",
    "- parameters.py: Contains all parameters needed for the files above.\n",
    "- utils.py: Utility functions.\n",
    "- kinematics.py: Contains inverse kinematic solver.\n",
    "\n",
    "In order to demonstrate the code in a more understandable way, let's go through the 3D reconstruction pipeline again but instead of clicking the buttons on the *Dash* we call the functions from this notebook directly.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### View sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View sampling\n",
    "views = dash.rc3dr.view_sampler.sample_views()\n",
    "num_views = len(views)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### View selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
   ],
   "source": [
    "# View selection - part 1\n",
    "\n",
    "# 1) TSDF to mesh\n",
    "F = dash.rc3dr.get_tsdf_volume()\n",
    "verts, faces = dash.rc3dr.view_selector.tsdf_to_mesh(F)\n",
    "\n",
    "# 2) Process mesh to get proper surface normals -> smooth mesh\n",
    "# NOTE w depends on mesh size\n",
    "w = 1e8\n",
    "v_star = utils.smooth_data(verts, faces, 1e8)\n",
    "\n",
    "# 3) Sample vertices from mesh weighted by curvature\n",
    "verts_choice, p = dash.rc3dr.view_selector.sample_verts_from_mesh(v_star, faces)\n",
    "normals = igl.per_vertex_normals(v_star, faces)\n",
    "# Note take unprocessed vertices\n",
    "verts_sam = verts[verts_choice,:]\n",
    "normals_sam = normals[verts_choice,:]\n",
    "\n",
    "plt.hist(p)\n",
    "\n",
    "ipv.figure()\n",
    "color = np.ones((v_star.shape[0], 3)) * np.expand_dims(p, axis=1) / np.max(p) + np.array([[0.5, 0,0.5]])\n",
    "mesh = ipv.plot_trisurf(v_star[:,0], v_star[:,1], v_star[:,2], triangles=faces, color=color)\n",
    "ipv.squarelim()\n",
    "ipv.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
   ],
   "source": [
    "# View selection - part 2\n",
    "\n",
    "# 4) Geometry values\n",
    "# Here we compute the scores for points of just one view\n",
    "view = [0,0, 0.75, 0, 1, 0, np.pi] # view[:3] = position, view[3:6] = rotation aixs, view[6] = angle\n",
    "# NOTE points outside view frustum get score = 0\n",
    "scores = dash.rc3dr.view_selector.compute_geometry_values(view, verts_sam, normals_sam)\n",
    "\n",
    "color = np.zeros((scores.shape[0],3))\n",
    "color[:,0] = scores\n",
    "\n",
    "plot_pts = np.vstack([verts_sam, np.array(view[:3])])\n",
    "plot_colors = np.vstack([color, np.array([1, 0, 0])])\n",
    "\n",
    "size = np.ones(plot_pts.shape[0])\n",
    "size[-1] = 2\n",
    "ipv.quickscatter(plot_pts[:,0], plot_pts[:,2], plot_pts[:,1], color = plot_colors, size=size, marker=\"sphere\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View selection - part 3\n",
    "\n",
    "# 5) Visibility values\n",
    "\n",
    "\n",
    "# The follwing function gets compiled to a cuda kernel. We use this function to track rays\n",
    "@cuda.jit\n",
    "def check_visibility_track_rays(track_rays, visibility, view_pos, points, tsdf_volume, K, voxel_size, mu):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        view: e_T_c (projection from camera view to end-effector)\n",
    "        points: points expressed in end-effector system. (n_points, 3)\n",
    "    \"\"\"\n",
    "    start_x = cuda.grid(1)\n",
    "    stride_x = cuda.gridsize(1)\n",
    "    \n",
    "    for i in range(start_x, points.shape[0], stride_x):\n",
    "        # ray marching from point to camera\n",
    "        # check if visible\n",
    "        \n",
    "        ray_point_x = points[i,0]\n",
    "        ray_point_y = points[i,1]\n",
    "        ray_point_z = points[i,2]\n",
    "        \n",
    "        ray_direction_x = view_pos[0] - points[i,0]\n",
    "        ray_direction_y = view_pos[1] - points[i,1]\n",
    "        ray_direction_z = view_pos[2] - points[i,2]\n",
    "        \n",
    "        # normalize ray direction\n",
    "        ray_norm = math.sqrt(ray_direction_x**2 + ray_direction_y**2 + ray_direction_z**2)\n",
    "        ray_direction_x /= ray_norm\n",
    "        ray_direction_y /= ray_norm\n",
    "        ray_direction_z /= ray_norm\n",
    "        \n",
    "        # NOTE Assumes quadratic volume\n",
    "        tsdf_dim = voxel_size * tsdf_volume.shape[0]\n",
    "        \n",
    "        while True:\n",
    "            # go in ray direction\n",
    "            ray_point_x += mu * ray_direction_x\n",
    "            ray_point_y += mu * ray_direction_y\n",
    "            ray_point_z += mu * ray_direction_z\n",
    "            \n",
    "            # get volume coordinates\n",
    "            tmp = ((ray_point_x + tsdf_dim / 2) / tsdf_dim) * tsdf_volume.shape[0]\n",
    "            v_x = int(math.floor(tmp))\n",
    "            tmp = ((ray_point_y + tsdf_dim / 2) / tsdf_dim) * tsdf_volume.shape[0]\n",
    "            v_y = int(math.floor(tmp))\n",
    "            tmp = ((ray_point_z + tsdf_dim / 2) / tsdf_dim) * tsdf_volume.shape[0]\n",
    "            v_z = int(math.floor(tmp))\n",
    "            \n",
    "            if (v_x < 0 or v_x >= tsdf_volume.shape[0] \n",
    "                or v_y < 0 or v_y >= tsdf_volume.shape[1] \n",
    "                or v_z < 0 or v_z >= tsdf_volume.shape[2]):\n",
    "                visibility[i] = 1\n",
    "                break\n",
    "                \n",
    "            tsdf_value = tsdf_volume[v_x, v_y, v_z]\n",
    "            track_rays[v_x, v_y, v_z] = 1\n",
    "            \n",
    "            if tsdf_value <= 0:\n",
    "                visibility[i] = 0\n",
    "                break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "visibility = np.zeros(verts_sam.shape[0])\n",
    "track_rays = np.zeros(F.shape)\n",
    "mu = rc3dr.tsdf_vol_params.voxel_size * np.sqrt(3)\n",
    "#execute kernel with 30 block each 128 threas\n",
    "check_visibility_track_rays[30,128](track_rays, \n",
    "                                    visibility, \n",
    "                                    np.array(view[:3]), \n",
    "                                    verts_sam, \n",
    "                                    F, \n",
    "                                    rc3dr.cam_params.K, \n",
    "                                    rc3dr.tsdf_vol_params.voxel_size, \n",
    "                                    mu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
   ],
   "source": [
    "# visualize rays\n",
    "ipv.quickvolshow(track_rays, level=1, opacity=1, level_width=0.1, data_min=0, data_max= 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
   ],
   "source": [
    "# visual visible points\n",
    "\n",
    "color = np.zeros((visibility.shape[0],3))\n",
    "color[:,0] = visibility\n",
    "plot_colors = np.vstack([color, np.array([1, 0, 0])])\n",
    "\n",
    "ipv.quickscatter(plot_pts[:,0], plot_pts[:,2], plot_pts[:,1], color = plot_colors, size=size, marker=\"sphere\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
   ],
   "source": [
    "# View selection - part 4\n",
    "\n",
    "# 6) Greedy selection\n",
    "scores = np.ndarray((num_views, verts_sam.shape[0]))\n",
    "for i in range(num_views):\n",
    "    scores[i,:] = dash.rc3dr.view_selector.compute_geometry_values(views[i], verts_sam, normals_sam)\n",
    "    visibility = dash.rc3dr.view_selector.compute_visibility_values(views[i], verts_sam, F)\n",
    "    scores[i,:] *= visibility\n",
    "    \n",
    "best_views, total_score = dash.rc3dr.view_selector.greedy_best_subset(scores)\n",
    "\n",
    "print(\"Total score: %f\" % (total_score))\n",
    "print(\"Best views: \\n\", best_views)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
   ],
   "source": [
    "# Visualization\n",
    "f_T_e = np.linalg.inv(rc3dr.tsdf_vol_params.e_T_f)\n",
    "views_f = utils.views_e_to_f(views, f_T_e)\n",
    "cam_poses = utils.views_to_cam_poses(views_f)\n",
    "\n",
    "ipv.figure()\n",
    "quiver = ipv.quiver(cam_poses[best_views,0], cam_poses[best_views,1], cam_poses[best_views,2], cam_poses[best_views,3], cam_poses[best_views,4], cam_poses[best_views,5], size=6, color=(0,1,0))\n",
    "quiver = ipv.quiver(cam_poses[:,0], cam_poses[:,1], cam_poses[:,2], cam_poses[:,3], cam_poses[:,4], cam_poses[:,5], size=5)\n",
    "\n",
    "scatter = ipv.scatter(verts_sam[:,0], verts_sam[:,1], verts_sam[:,2], size=1, marker=\"sphere\")\n",
    "ipv.squarelim()\n",
    "ipv.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### View traversal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "view_traverser = view_traversal.ViewTraverser(rc3dr.view_traverser_params, \n",
    "                                              rc3dr.cam_params, \n",
    "                                              rc3dr.tsdf_vol_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
   ],
   "source": [
    "# NOTE Multiple users can connect to the dash\n",
    "# For each user we save data such as TSDF/ robot/ etc.\n",
    "user_sid = list(dash.rc3dr.user_sids)[-1]\n",
    "print(user_sid)\n",
    "\n",
    "robot = dash.rc3dr.robot_map[user_sid]\n",
    "ee_id = robot.index(dash.rc3dr.ee_name)\n",
    "inv_kin_solver = dash.rc3dr.inv_kin_solver\n",
    "\n",
    "views_best = np.array(views)[best_views,:].tolist()\n",
    "\n",
    "q_steps, greedy_traversal, view_steps_e = view_traverser.get_steps(views_best, robot, ee_id, inv_kin_solver)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "view_steps_f = utils.views_e_to_f(view_steps_e, f_T_e)\n",
    "cam_pose_steps = utils.views_to_cam_poses(view_steps_f)\n",
    "\n",
    "ipv.figure()\n",
    "quiver = ipv.quiver(cam_poses[best_views,0], cam_poses[best_views,1], cam_poses[best_views,2], cam_poses[best_views,3], cam_poses[best_views,4], cam_poses[best_views,5], size=10, color=(0,1,0))\n",
    "quiver = ipv.quiver(cam_pose_steps[:,0], cam_pose_steps[:,1], cam_pose_steps[:,2], cam_pose_steps[:,3], cam_pose_steps[:,4], cam_pose_steps[:,5], size=5)\n",
    "\n",
    "scatter = ipv.scatter(verts_sam[:,0], verts_sam[:,1], verts_sam[:,2], size=1, marker=\"sphere\")\n",
    "ipv.squarelim()\n",
    "ipv.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TODOs\n",
    "\n",
    "- Convert 8-bit depth channel to RGB or single float channel. (At the moment the depth channel can only take 255 values. However, the depth channel has higher resolution at near distance than at far distance).\n",
    "- Improve view sampling\n",
    "- Improve path trajectory planning\n",
    "- Implement framework in C++\n",
    "- ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusions\n",
    "\n",
    "This project presents a complete solution for the robot-guided 3D reconstruction (RC3DR) process in simulation. In a next step, we want to evaluate this framework in the real world."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.5 64-bit ('tensorflow_2': conda)",
   "language": "python",
   "name": "python37564bittensorflow2conda7920346893e74f3a9f64031d682baec0"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}