Merge pull request andyzeng#13 from kevinzakka/cpu-optim

added numba optimisation to cpu code. ~2x FPS speedup.

Author: andyzeng

.gitignore

@@ -1 +1,2 @@
 __pycache__
+*.ply

README.md

@@ -8,9 +8,9 @@ An older CUDA/C++ version can be found [here](https://github.com/andyzeng/tsdf-f
 
 ## Requirements
 
-* Python 2.7+ with [NumPy](http://www.numpy.org/), [PyCUDA](https://developer.nvidia.com/pycuda), [OpenCV](https://docs.opencv.org/3.0-beta/doc/py_tutorials/py_tutorials.html), [Scikit-image](https://scikit-image.org/). These can be quickly installed/updated by running the following:
+* Python 2.7+ with [NumPy](http://www.numpy.org/), [PyCUDA](https://developer.nvidia.com/pycuda), [OpenCV](https://docs.opencv.org/3.0-beta/doc/py_tutorials/py_tutorials.html), [Scikit-image](https://scikit-image.org/) and [Numba](https://numba.pydata.org/). These can be quickly installed/updated by running the following:
   ```shell
-  pip install --user numpy opencv-python scikit-image
+  pip install --user numpy opencv-python scikit-image numba
   ```
 * [Optional] GPU acceleration requires an NVIDA GPU with [CUDA](https://developer.nvidia.com/cuda-downloads) and [PyCUDA](https://developer.nvidia.com/pycuda):
   ```shell
@@ -42,10 +42,10 @@ python demo.py
 This repository is a part of [3DMatch Toolbox](https://github.com/andyzeng/3dmatch-toolbox). If you find this code useful in your work, please consider citing:
 
 ```
-@inproceedings{zeng20163dmatch, 
-    title={3DMatch: Learning Local Geometric Descriptors from RGB-D Reconstructions}, 
-    author={Zeng, Andy and Song, Shuran and Nie{\ss}ner, Matthias and Fisher, Matthew and Xiao, Jianxiong and Funkhouser, Thomas}, 
-    booktitle={CVPR}, 
-    year={2017} 
+@inproceedings{zeng20163dmatch,
+    title={3DMatch: Learning Local Geometric Descriptors from RGB-D Reconstructions},
+    author={Zeng, Andy and Song, Shuran and Nie{\ss}ner, Matthias and Fisher, Matthew and Xiao, Jianxiong and Funkhouser, Thomas},
+    booktitle={CVPR},
+    year={2017}
 }
 ```

demo.py

@@ -1,52 +1,63 @@
-#!/usr/bin/env python
+"""Fuse 1000 RGB-D images from the 7-scenes dataset into a TSDF voxel volume with 2cm resolution.
+"""
 
-import numpy as np
-import cv2
 import time
-import fusion
 
+import cv2
+import numpy as np
+
+import fusion
 
-# (Optional) sample code to compute 3D bounds (in world coordinates) around convex hull of all camera view frustums in dataset
-print("Estimating voxel volume bounds...")
-n_imgs = 1000
-cam_intr = np.loadtxt("data/camera-intrinsics.txt",delimiter=' ')
-vol_bnds = np.zeros((3,2))
-for i in range(n_imgs):
 
+if __name__ == "__main__":
+  # ======================================================================================================== #
+  # (Optional) This is an example of how to compute the 3D bounds
+  # in world coordinates of the convex hull of all camera view
+  # frustums in the dataset
+  # ======================================================================================================== #
+  print("Estimating voxel volume bounds...")
+  n_imgs = 1000
+  cam_intr = np.loadtxt("data/camera-intrinsics.txt", delimiter=' ')
+  vol_bnds = np.zeros((3,2))
+  for i in range(n_imgs):
     # Read depth image and camera pose
-    depth_im = cv2.imread("data/frame-%06d.depth.png"%(i),-1).astype(float)/1000. # depth is saved in 16-bit PNG in millimeters
-    depth_im[depth_im == 65.535] = 0 # set invalid depth to 0 (specific to 7-scenes dataset)
-    cam_pose = np.loadtxt("data/frame-%06d.pose.txt"%(i)) # 4x4 rigid transformation matrix
+    depth_im = cv2.imread("data/frame-%06d.depth.png"%(i),-1).astype(float)
+    depth_im /= 1000.  # depth is saved in 16-bit PNG in millimeters
+    depth_im[depth_im == 65.535] = 0  # set invalid depth to 0 (specific to 7-scenes dataset)
+    cam_pose = np.loadtxt("data/frame-%06d.pose.txt"%(i))  # 4x4 rigid transformation matrix
 
     # Compute camera view frustum and extend convex hull
-    view_frust_pts = fusion.get_view_frustum(depth_im,cam_intr,cam_pose)
-    vol_bnds[:,0] = np.minimum(vol_bnds[:,0],np.amin(view_frust_pts,axis=1))
-    vol_bnds[:,1] = np.maximum(vol_bnds[:,1],np.amax(view_frust_pts,axis=1))
-
-# ---------------------------------------------------------------------
-
-# Initialize voxel volume
-print("Initializing voxel volume...")
-tsdf_vol = fusion.TSDFVolume(vol_bnds,voxel_size=0.02)
-
-# Loop through RGB-D images and fuse them together
-t0_elapse = time.time()
-for i in range(n_imgs):
-    print("Fusing frame %d/%d"%(i+1,n_imgs))
+    view_frust_pts = fusion.get_view_frustum(depth_im, cam_intr, cam_pose)
+    vol_bnds[:,0] = np.minimum(vol_bnds[:,0], np.amin(view_frust_pts, axis=1))
+    vol_bnds[:,1] = np.maximum(vol_bnds[:,1], np.amax(view_frust_pts, axis=1))
+  # ======================================================================================================== #
+
+  # ======================================================================================================== #
+  # Integrate
+  # ======================================================================================================== #
+  # Initialize voxel volume
+  print("Initializing voxel volume...")
+  tsdf_vol = fusion.TSDFVolume(vol_bnds, voxel_size=0.02)
+
+  # Loop through RGB-D images and fuse them together
+  t0_elapse = time.time()
+  for i in range(n_imgs):
+    print("Fusing frame %d/%d"%(i+1, n_imgs))
 
     # Read RGB-D image and camera pose
-    color_image = cv2.cvtColor(cv2.imread("data/frame-%06d.color.jpg"%(i)),cv2.COLOR_BGR2RGB)
-    depth_im = cv2.imread("data/frame-%06d.depth.png"%(i),-1).astype(float)/1000. # depth is saved in 16-bit PNG in millimeters
-    depth_im[depth_im == 65.535] = 0 # set invalid depth to 0 (specific to 7-scenes dataset)
-    cam_pose = np.loadtxt("data/frame-%06d.pose.txt"%(i)) # 4x4 rigid transformation matrix
+    color_image = cv2.cvtColor(cv2.imread("data/frame-%06d.color.jpg"%(i)), cv2.COLOR_BGR2RGB)
+    depth_im = cv2.imread("data/frame-%06d.depth.png"%(i),-1).astype(float)
+    depth_im /= 1000.
+    depth_im[depth_im == 65.535] = 0
+    cam_pose = np.loadtxt("data/frame-%06d.pose.txt"%(i))
 
     # Integrate observation into voxel volume (assume color aligned with depth)
-    tsdf_vol.integrate(color_image,depth_im,cam_intr,cam_pose,obs_weight=1.)
+    tsdf_vol.integrate(color_image, depth_im, cam_intr, cam_pose, obs_weight=1.)
 
-fps = n_imgs/(time.time()-t0_elapse)
-print("Average FPS: %.2f"%(fps))
+  fps = n_imgs / (time.time() - t0_elapse)
+  print("Average FPS: {:.2f}".format(fps))
 
-# Get mesh from voxel volume and save to disk (can be viewed with Meshlab)
-print("Saving to mesh.ply...")
-verts,faces,norms,colors = tsdf_vol.get_mesh()
-fusion.meshwrite("mesh.ply",verts,faces,norms,colors)
+  # Get mesh from voxel volume and save to disk (can be viewed with Meshlab)
+  print("Saving to mesh.ply...")
+  verts, faces, norms, colors = tsdf_vol.get_mesh()
+  fusion.meshwrite("mesh.ply", verts, faces, norms, colors)