小白科研笔记：从零学习基于TSDF的三维重建

1. 前言

这篇博客主要分析这个老哥做的Python实现的基于TSDF地图的三维重建，点击它获取Github链接。三维重建的效果图如下所示。这篇博客主要分析两件事情，一是TSDF地图的原理和融合理论，二是TSDF实现的代码。如果后续有时间，我再去做一个mini版本的三维重建小工程。有关Kinect Fusion相关的知识可以参考这篇知乎笔记。

图1：三维重建效果图

2. TSDF地图的构成和初始化

TSDF地图是由一堆体素构成的，每个体素内包含两个变量，一是tsdf值（用于生成重建表面），二是RGB信息（给重建表面贴上彩色纹理）。TSDF地图与实际场景的物理坐标系是可转换的。这是这一节讲解的基本信息。tsdf值和rgb值的计算会在下一节讨论。

相信大家即便是没有玩过“我的世界”，也肯定听说过这款游戏。在这款游戏中，世界是由三维小格子组成的。TSDF地图的原理就和它很像。TSDF是截断符号函数（Truncated Signed Distance Function）的缩写（初学不必去理会这些英文字眼）。一个三维的TSDF地图就是由$L\times W\times H$个三维小方块组成。我们称三维小方块为体素（Voxel）。列举一个例子。给定一个尺寸为$100\times 100\times 100$的TSDF地图，该地图是有尺寸为$0.05m\times 0.05m\times 0.05m$的体素构成。那么可以计算出该TSDF地图覆盖一个$5m\times 5m\times 5m$的真实三维场景。这个地图中有$10^6$个体素。每一个体素记录着tsdf值和RGB颜色信息（tsdf值范围是$[-1,1]$，在第三节做介绍）。GPU需要大概2KB存储单个体素的所有信息。因此GPU需要大概2GB左右的显存空间存放这个TSDF地图。从这些计算过程可见，TSDF占据非常大的显存空间，它适合小场景下的三维重建，比如室内环境。

接下来定量地讨论一下TSDF地图下体素坐标系和物理坐标系的变换问题。我们知道TSDF地图是一个大立方块，里面是由很多小立方块（体素）构成的。与此同时，TSDF地图对应着实际的立方块空间。记TSDF地图中$(0,0,0)$处体素对应实际场景坐标为$(x_0,y_0,z_0)$点。那么场景中坐标为$(x,y,z)$的点，对应TSDF地图中$\text{int}((x-x_0)/\text{Voxel.x},(y-y_0)/\text{Voxel.y},(z-z_0)/\text{Voxel.z})$处的体素。因为体素坐标系是整数，所以实际点坐标在转化为TSDF地图下体素坐标系的时候需要取整。

TSDF地图由体素构成，每一个体素包含一个$[-1,1]$的tsdf值，以及包含一个rgb值。因此初始化一个TSDF地图，相当于初始化两个$[0,L-1]\times [0,W-1]\times [0,H-1]$范围内的三维张量，分别存放tsdf值和rgb值。使用numpy就能轻松搞定。把TSDF地图从CPU存储中移动到GPU中需要两步，第一步是调用cuda.mem_alloc在GPU中提前开辟一块内存空间，第二步是调用cuda.memcpy_htod，用这块内存空间存储TSDF地图数据。htod表示host to device，即从CPU存数据到GPU。如果是反过来的话，就使用cuda.memcpy_dtoh。注意，在初始化的时候，TSDF地图内所有体素的tsdf值被填充为1。这个原因会在下一节讲解。

# self._vol_dim 表示 TSDF 地图的尺寸，[L,W,H]
# 使用 np.ones 后，_tsdf_vol_cpu 是一个 L*W*H 的张量，表示 TSDF 地图，
# 填充值为 tsdf值，在初始化阶段，默认为 1
self._tsdf_vol_cpu = np.ones(self._vol_dim).astype(np.float32)
# 对应颜色信息，在初始化阶段，默认为 0
self._color_vol_cpu = np.zeros(self._vol_dim).astype(np.float32)
# 简而言之，TSDF 地图就由 tsdf 值张量和颜色张量组成。
# 把 TSDF 地图存放在 GPU 中
self._tsdf_vol_gpu = cuda.mem_alloc(self._tsdf_vol_cpu.nbytes)
cuda.memcpy_htod(self._tsdf_vol_gpu,self._tsdf_vol_cpu)
self._color_vol_gpu = cuda.mem_alloc(self._color_vol_cpu.nbytes)
cuda.memcpy_htod(self._color_vol_gpu,self._color_vol_cpu)

3. 如何计算一个体素的tsdf值和rgb值

这一节讨论这样一件事情，在TSDF地图中随便选出一个体素$x$，怎样去计算它的tsdf值和rgb值呢？计算它们的已知条件是需要一张深度图。首先去讲解tsdf值的计算。tsdf值的前身是sdf值。咱们先看下面的经典讲解图。白灰色的小方格表示TSDF地图中的各个体素。蓝色的三角形表示相机的视场范围。图中间有一条绿色的截线，表示一个物体的截面。深度相机的深度图可以显示物体截面的深度信息。我们要去计算体素$x$的sdf值和tsdf值，即$\text{sdf}(x)$和$\text{tsdf}(x)$。

图2：讲解tsdf值和sdf值的示意图

首先，需要计算体素$x$在物理坐标下的位置。记体素$x$在TSDF地图上的坐标是$(v_x,v_y,v_z)$。那么该体素在物理世界坐标系下的位置是$P_{x,wrd}=(x_0+v_x\cdot \text{Voxel}.x,y_0+v_y\cdot \text{Voxel}.y,z_0+v_z\cdot \text{Voxel}.z)$。然后需要计算体素$x$在相机坐标系下的位置。设相机相对于物理坐标系下的位姿是$\text{R}$和$T$。那么体素$x$在相机坐标系下的位置是$P_{x,cam}=\text{R}{P}_{x,wrd}+T$。根据相机成像模型，$ca{m}_z(x)\cdot I_x=\text{K}{P}_{x,cam}$。$\text{K}$表示相机的内参数矩阵。$I_x$表示体素$x$投影在相机成像平面下的像素坐标。$ca{m}_z(x)$表示体素$x$相对于相机的深度。

沿着相机的光心和体素$x$做一条直线（图中显示为深蓝色的粗线），这条线会与物体的截面有一个交点，这个交点记为$P$点。$P$点的深度记为$d_P$。在实际计算中，$d_P$的获取真的要做直线，计算交点有这么麻烦吗？完全不是，那只是个示意图讲解而已。记当前的深度图为$D$。在实际计算中$d_p=D(I_x)$。那么体素$x$的sdf值就可以计算出来了：

$\text{sdf}(x)=d_p-d_x=D(I_x)-ca{m}_z(x)$

可以发现：$\text{sdf}(x)>0$表示体素$x$处于相机和物体表面之间；$\text{sdf}(x)<0$表示体素$x$处于物体表面之后。根据体素$x$的sdf值，可以计算它的的tsdf值：

$\text{tsdf}(x)=\max [-1,\min (1,\text{sdf}(x)/t)]$

tsdf的计算表达式看上去有些古怪。首先是要做一个除法操作，即$\text{sdf}(x)/t$，这就是截断（Truncated）的含义。当$-t\le \text{sdf}(x)\le t$时，$\text{tsdf}(x)=\text{sdf}(x)/t\in [-1,1]$。当$\text{sdf}(x)>t$或者$\text{sdf}(x)<-t$时，$\text{tsdf}(x)=1$或者$-1$。其实这个表达式是有物理意义的。$t$可以看作是体素$x$和截面对应点$P$深度差值的阈值。当体素$x$距离截面对应点$P$非常远的时候，它的tsdf值等于正负一。当体素$x$距离截面对应点$P$比较近的时候，它的tsdf值$[-1,1]$之间，是有意义的。用更为通俗的话说，当体素离表面非常近的时候，它的tsdf值接近于零；当体素离表面非常远的时候，它的tsdf值趋于正一或者负一。这正好对应于图2。

还记不记得，在初始化的时候，我们设所有体素的tsdf值为1，相当于这个TSDF地图中没有任何表面。所以这个初始化是有意义的。

最后再去讨论体素$x$的rgb值计算方法。其实很简单，记深度图对应的RGB图像为$\text{RGB}$。那么$\text{rgb}(x)=\text{RGB}(I_x)$。实际代码实现的时候会稍稍跟它不太一样，这个放在后面去讲。对每一个体素，都是需要重复进行上述的过程，所以代码实现利用并行计算GPU会非常高效。这一节的计算代码我会放在第四节去讲。

4. TSDF地图融合

在第三节，咱们讨论“给单张深度图和RGB图像，计算各个体素的tsdf值和rgb值”的过程。这个过程用数学语言描述就是：$\text{TSDF}(t)=f(D(t),\text{RGB}(t),\text{R}(t),T(t))$。$\text{R}(t),T(t)$表示在$T=t$时刻下相机相对于物理坐标系下的位姿。$\text{TSDF}(t)$由两个三维张量组成，分别是$\text{tsdf}(t,x)$和$\text{rgb}(t,x)$。整型三维向量$x$是TSDF地图中的体素索引。算子$f(\cdot )$就是第三节的计算过程。那么对于一个完整的RGB-D图序列，怎样去处理呢？势必希望设计一个算子$g(\cdot )$，使得$\text{Result\_tsdf}=g(\text{tsdf}(0),\text{tsdf}(1),...,\text{tsdf}(N))$，把各个时刻计算的tsdf值做个融合。

其实$\text{Result\_tsdf}=g(\text{tsdf}(0),\text{tsdf}(1),...,\text{tsdf}(N))$是离线三维重建的表述，因为它把之前时刻的结果一古脑地塞进$g(\cdot )$中，并不适用于实时三维重建。实时三维重建的数学表达式应该是这样的：

$\text{tsdf}(0)=\text{tsdf}(0)$
$\text{tsdf}(1)=g(\text{tsdf}(0),\text{tsdf}(1))$

写成一般表达式就是：

$\text{tsdf}(t)=g(\text{tsdf}(t-1),\text{tsdf}(t)),t\ge 1$

在Kinect Fusion甚至是更早期的三维重建算法中，把$g(\cdot )$设计为一种加权平均的形式：

$\text{tsdf}(t)=[W(t-1)\cdot \text{tsdf}(t-1)+w(t)\cdot \text{tsdf}(t)]/[W(t-1)+w(t)]$

$W(t)=W(t-1)+w(t)$

$W(0)={\text{0}}_{L\times W\times H}$

$w(t)={\text{1}}_{L\times W\times H}$

注意$W(t)$和$w(t)$都是三维张量，表示的都是TSDF中每一个体素的权值。上述公式中的操作符"$\cdot$"，"$+$"，“/”均表示张量间点对点的相乘相加和相除。$W(0)$初始化为零值张量。$w(t)$则是常数值张量。为什么$w(t)$会设计为常数值张量呢？细品一下，这是有一种时间累计效应的感觉。$\text{tsdf}(t-1)$累计了$t\in [0,t-1]$时刻所有的三维重建结果，它的重要性自然比$\text{tsdf}(t)$高出很多。因为$W(t)$也是三维张量，所以也要放在一开始做个初始化哈。

    self._weight_vol_cpu = np.zeros(self._vol_dim).astype(np.float32)
    self._weight_vol_gpu = cuda.mem_alloc(self._weight_vol_cpu.nbytes)
    cuda.memcpy_htod(self._weight_vol_gpu,self._weight_vol_cpu)

此外，从$\text{TSDF}(t)=g(\text{TSDF}(t-1),\text{TSDF}(t)),t\ge 1$这个公式看出，体素的tsdf值的计算和tsdf值的融合可以放在一块进行。下面会缓缓给出TSDF融合的代码，在这个Python实现的TSDF的三维重建中，这段代码是用Cuda写的，通过PyCuda的方式编译这段Cuda代码，让Python主程序能够调用这个Cuda API。

对了，对于RGB信息的融合，也是采取了类似的形式：

$\text{rgb}(t)=g(\text{rgb}(t-1),\text{rgb}(t)),t\ge 1$

$\text{rgb}(t)=[W(t-1)\cdot \text{rgb}(t-1)+w(t)\cdot \text{rgb}(t)]/[W(t-1)+w(t)]$

上述公式中的操作符"$\cdot$"，"$+$"，“/”同样表示张量间点对点的相乘相加和相除。

编程细节一：怎样选取线程块Block和线程网格Grid？

在体素的tsdf值的计算和tsdf值的融合的过程中，需要遍历TSDF中每一个体素。对于一个$L\times W\times H$的TSDF地图，就自然要遍历$N=LWH$数量个体素，数量级大概在GB级或者说是百万级。GPU虽然是多线程并行运算的利器，但是也不可能同时开百万级的线程。GPU一次性可以开几十万个线程，数量记为$M$。简单地讲，它需要重复地运作$N/M$次（流水线操作可以加速这个过程，我是指微机原理中处理指令的那个流水线），可以遍历完全部的体素。这里不去细述GPU的更深的原理（我自己也不是太懂）。总之，GPU需要用到线程块Block和线程网格Grid完成对线程的调度。

对GPU来说，Block和Grid的选取非常重要。只有选取对了，GPU才能遍历完TSDF中全部的体素，于此同时不会浪费太多的多余线程。先看看Python中调用tsdf融合的_cuda_integrate的代码。这里有关PyCuda的知识不去做介绍。咱们比较关注代码中block和grid的设定。

        self._cuda_integrate = self._cuda_src_mod.get_function("integrate")
        self._cuda_integrate(self._tsdf_vol_gpu,
                            self._weight_vol_gpu,
                            self._color_vol_gpu,
                            cuda.InOut(self._vol_dim.astype(np.float32)),
                            cuda.InOut(self._vol_origin.astype(np.float32)),
                            cuda.InOut(cam_intr.reshape(-1).astype(np.float32)),
                            cuda.InOut(cam_pose.reshape(-1).astype(np.float32)),
                            cuda.InOut(np.asarray([
                              gpu_loop_idx,
                              self._voxel_size,
                              im_h,
                              im_w,
                              self._trunc_margin,
                              obs_weight
                            ], np.float32)),
                            cuda.InOut(color_im.reshape(-1).astype(np.float32)),
                            cuda.InOut(depth_im.reshape(-1).astype(np.float32)),
                            block=(self._max_gpu_threads_per_block,1,1),
                            grid=(
                              int(self._max_gpu_grid_dim[0]),
                              int(self._max_gpu_grid_dim[1]),
                              int(self._max_gpu_grid_dim[2]),
                            )

上述代码中，block=(self._max_gpu_threads_per_block,1,1)定义了一个线程块。block是一个三元组，其中block.x=_max_gpu_threads_per_block，block.y=1，block.z=1。这个线程块可以一次性地打开block.x*block.y*block.z个线程。grid表示线程网格，它也是一个三元组。线程网格表示它来调度grid.x*grid.y*grid.z个线程块。我们可以看到一个上下级的关系：线程网格调度线程块，线程块则控制线程。为了让GPU能够遍历到所有的体素，线程网格和线程块的设计必须满足下面的式子：

$\text{block}.x\cdot \text{block}.y\cdot \text{block}.z\cdot \text{grid}.x\cdot \text{grid}.y\cdot \text{grid}.z\ge L\cdot W\cdot H$

说句题外话。细心一点可以注意到，block和grid的设定就是一个排列组合的问题，似乎block和grid的选型方案有很多组。但是在实际应用中，block和grid设置地过高和过低都会对并行计算产生影响，甚至GPU会罢工，这个原因可以追溯到GPU的硬件设计上。有兴趣可以看看GPU高性能编程相关书籍。在这里，我们就随便的设计一下就行啦，只要满足上面的式子即可。所以block和grid的设定代码如下所示：

      self._max_gpu_threads_per_block = gpu_dev.MAX_THREADS_PER_BLOCK
      n_blocks = int(np.ceil(float(np.prod(self._vol_dim))/float(self._max_gpu_threads_per_block)))
      grid_dim_x = min(gpu_dev.MAX_GRID_DIM_X,int(np.floor(np.cbrt(n_blocks))))
      grid_dim_y = min(gpu_dev.MAX_GRID_DIM_Y,int(np.floor(np.sqrt(n_blocks/grid_dim_x))))
      grid_dim_z = min(gpu_dev.MAX_GRID_DIM_Z,int(np.ceil(float(n_blocks)/float(grid_dim_x*grid_dim_y))))
      self._max_gpu_grid_dim = np.array([grid_dim_x,grid_dim_y,grid_dim_z]).astype(int)

假设一个极端地情况，grid_dim_x到达MAX_GRID_DIM_X，grid_dim_y到达MAX_GRID_DIM_Y，grid_dim_z到达MAX_GRID_DIM_Z，但是线程数还是小于体素的数量，这该怎么办呢？即如下的数学表达式：

$M=\text{block}.x\cdot \text{block}.y\cdot \text{block}.z\cdot \text{grid\_max}.x\cdot \text{grid\_max}.y\cdot \text{grid\_max}.z<L\cdot W\cdot H$

一个简单的办法就是多次遍历，我遍历个$[LWH/M]$次就够啦！

# 一般情况下，如果线程够用，_n_gpu_loops 就等于 1
self._n_gpu_loops = int(np.ceil(float(np.prod(self._vol_dim))/float(np.prod(self._max_gpu_grid_dim)*self._max_gpu_threads_per_block)))

编程细节二：如何根据线程块和线程网格遍历体素？

前文已经讲了GPU中的上下级的关系：线程网格调度线程块，线程块则控制线程。直接上代码，就不难理解了，可以这样得到体素的索引voxel_idx：

  int gpu_loop_idx = (int) other_params[0]; # gpu_loop_idx = 1
  int max_threads_per_block = blockDim.x;
  int block_idx = blockIdx.z*gridDim.y*gridDim.x+blockIdx.y*gridDim.x+blockIdx.x;
  int voxel_idx = gpu_loop_idx*gridDim.x*gridDim.y*gridDim.z*max_threads_per_block+block_idx*max_threads_per_block+threadIdx.x;

然而体素的三维张量在GPU中存储是一维张量。所以要想得到该体素的体素坐标，需要如下的操作：

    // Get voxel grid coordinates (note: be careful when casting)
    float voxel_x = floorf(((float)voxel_idx)/((float)(vol_dim_y*vol_dim_z)));
    float voxel_y = floorf(((float)(voxel_idx-((int)voxel_x)*vol_dim_y*vol_dim_z))/((float)vol_dim_z));
    float voxel_z = (float)(voxel_idx-((int)voxel_x)*vol_dim_y*vol_dim_z-((int)voxel_y)*vol_dim_z);

理解了上述两个编程细节，剩下的过程就很好理解了：

      # Cuda kernel function (C++)
      self._cuda_src_mod = SourceModule("""
        __global__ void integrate(float * tsdf_vol,
                                  float * weight_vol,
                                  float * color_vol,
                                  float * vol_dim,
                                  float * vol_origin,
                                  float * cam_intr,
                                  float * cam_pose,
                                  float * other_params,
                                  float * color_im,
                                  float * depth_im) {
          // Get voxel index
          int gpu_loop_idx = (int) other_params[0];
          int max_threads_per_block = blockDim.x;
          int block_idx = blockIdx.z*gridDim.y*gridDim.x+blockIdx.y*gridDim.x+blockIdx.x;
          int voxel_idx = gpu_loop_idx*gridDim.x*gridDim.y*gridDim.z*max_threads_per_block+block_idx*max_threads_per_block+threadIdx.x;
          int vol_dim_x = (int) vol_dim[0];
          int vol_dim_y = (int) vol_dim[1];
          int vol_dim_z = (int) vol_dim[2];
          if (voxel_idx > vol_dim_x*vol_dim_y*vol_dim_z)
              return;
          // Get voxel grid coordinates (note: be careful when casting)
          float voxel_x = floorf(((float)voxel_idx)/((float)(vol_dim_y*vol_dim_z)));
          float voxel_y = floorf(((float)(voxel_idx-((int)voxel_x)*vol_dim_y*vol_dim_z))/((float)vol_dim_z));
          float voxel_z = (float)(voxel_idx-((int)voxel_x)*vol_dim_y*vol_dim_z-((int)voxel_y)*vol_dim_z);
          // Voxel grid coordinates to world coordinates
          float voxel_size = other_params[1];
          float pt_x = vol_origin[0]+voxel_x*voxel_size;
          float pt_y = vol_origin[1]+voxel_y*voxel_size;
          float pt_z = vol_origin[2]+voxel_z*voxel_size;
          // World coordinates to camera coordinates
          float tmp_pt_x = pt_x-cam_pose[0*4+3];
          float tmp_pt_y = pt_y-cam_pose[1*4+3];
          float tmp_pt_z = pt_z-cam_pose[2*4+3];
          float cam_pt_x = cam_pose[0*4+0]*tmp_pt_x+cam_pose[1*4+0]*tmp_pt_y+cam_pose[2*4+0]*tmp_pt_z;
          float cam_pt_y = cam_pose[0*4+1]*tmp_pt_x+cam_pose[1*4+1]*tmp_pt_y+cam_pose[2*4+1]*tmp_pt_z;
          float cam_pt_z = cam_pose[0*4+2]*tmp_pt_x+cam_pose[1*4+2]*tmp_pt_y+cam_pose[2*4+2]*tmp_pt_z;
          // Camera coordinates to image pixels
          int pixel_x = (int) roundf(cam_intr[0*3+0]*(cam_pt_x/cam_pt_z)+cam_intr[0*3+2]);
          int pixel_y = (int) roundf(cam_intr[1*3+1]*(cam_pt_y/cam_pt_z)+cam_intr[1*3+2]);
          // Skip if outside view frustum
          int im_h = (int) other_params[2];
          int im_w = (int) other_params[3];
          if (pixel_x < 0 || pixel_x >= im_w || pixel_y < 0 || pixel_y >= im_h || cam_pt_z<0)
              return;
          // Skip invalid depth
          float depth_value = depth_im[pixel_y*im_w+pixel_x];
          if (depth_value == 0)
              return;
          // Integrate TSDF
          float trunc_margin = other_params[4];
          float depth_diff = depth_value-cam_pt_z;
          if (depth_diff < -trunc_margin)
              return;
          float dist = fmin(1.0f,depth_diff/trunc_margin);
          float w_old = weight_vol[voxel_idx];
          float obs_weight = other_params[5];
          float w_new = w_old + obs_weight;
          weight_vol[voxel_idx] = w_new;
          tsdf_vol[voxel_idx] = (tsdf_vol[voxel_idx]*w_old+obs_weight*dist)/w_new;
          // Integrate color
          float old_color = color_vol[voxel_idx];
          float old_b = floorf(old_color/(256*256));
          float old_g = floorf((old_color-old_b*256*256)/256);
          float old_r = old_color-old_b*256*256-old_g*256;
          float new_color = color_im[pixel_y*im_w+pixel_x];
          float new_b = floorf(new_color/(256*256));
          float new_g = floorf((new_color-new_b*256*256)/256);
          float new_r = new_color-new_b*256*256-new_g*256;
          new_b = fmin(roundf((old_b*w_old+obs_weight*new_b)/w_new),255.0f);
          new_g = fmin(roundf((old_g*w_old+obs_weight*new_g)/w_new),255.0f);
          new_r = fmin(roundf((old_r*w_old+obs_weight*new_r)/w_new),255.0f);
          color_vol[voxel_idx] = new_b*256*256+new_g*256+new_r;
        }""")

5. 结束语

$\text{R}(t),T(t)$需要用到ICP去计算，在这个Github开源项目中没有去写，而是利用了数据集的真值。不管怎样，这个小哥写的TSDF地图融合代码非常清楚，值得一读 😃

最后祝大家青年节快乐！