【NeRF】基于Mindspore的NeRF实现
一、NeRF介绍
1. 背景
传统计算机图形学技术经过几十年发展，主要技术路线已经相对稳定。随着深度学习技术的发展，新兴的神经渲染技术给计算机图形学带来了新的机遇，受到了学界和工业界的广泛关注。神经渲染是深度网络合成图像的各类方法的总称，各类神经渲染的目标是实现图形渲染中建模和渲染的全部或部分的功能，基于神经辐射场（Neural Radiance Field, NeRF）的场景三维重建是近期神经渲染的热点方向，目标是使用神经网络实现新视角下的2D图像生成，在20和21年的CVPR、NeuIPS等各大AI顶会上，我们可以看到几十、上百篇相关的高水平论文。
2. 网络结构
NeRF使用多层感知机（Multilayer Perceptron，MLP）来重建三维场景，也就是去拟合空间点的颜色分布和光密度分布的函数。NeRF的网络结构如图1所示，网络的输入是采样点对应的空间坐标和视角，输出是采样点对应的密度和RGB值。
由于颜色\boldsymbol{c}c和光密度\sigmaσ在空间中并不是平滑的，变化是比较剧烈的，这意味着函数存在很多高频的部分，让模型去表示这种函数比较困难，所以NeRF通过positional encoding，对输入的\boldsymbol{r},\boldsymbol{d}r,d进行编码、升维，从而能够让模型更好地学出场景的一些细节部分，具体映射方式如下所示，该映射将标量pp映射成一个2L+12L+1维的向量：
\boldsymbol{\gamma}(p)=\left[p,\sin(2^0\pi p),\cos(2^0\pi p),\cdots,\sin(2^{L-1}\pi p), \cos(2^{L-1}\pi p)\right]γ(p)=[p,sin(20πp),cos(20πp),⋯,sin(2L−1πp),cos(2L−1πp)]
NeRF采用的MLP完整架构如下图所示：
3. 光线步进法
NeRF使用MLP隐式重建三维场景时的输入是采样点的位姿，NeRF的目标是实现2D新视角图像生成，那么要怎么得到采样点的位姿，又怎么使用重建得到的3D密度和颜色呢，得到新视角下的2D图像？
为了处理这两个问题，NeRF使用了光线步进（Ray Marching）这一经典方法，设\boldsymbol{o}o代表相机原点O_cOc在世界坐标系中的位置，\boldsymbol{d}d代表射线的单位方向矢量，tt代表从O_cOc出发，沿射线方向行进的距离。使用光线步进法时，2D图像的每个像素对应于一条射线，每条射线上的任意位置可以表示为\boldsymbol{r}(t)=\boldsymbol{o}+t\boldsymbol{d}r(t)=o+td，对这些射线进行采样（具体使用下文介绍的随机采样和重要性采样）即可得到采样点的位姿，对这些射线上的采样点进行积分（具体使用下文介绍的体绘制方法）即可得到2D像素的RGB值。
4. 随机采样和重要性采样
NeRF使用随机采样和重要性采样结合的方式在光线步进法生成的射线上进行采样，这是因为空间中的物体分布是稀疏的，一条射线上可能只有很小的一段区域是对最终渲染起作用的，如果用均匀采样会浪费很多采样点，网络也难以学到整个连续空间中的分布，所以采用coarse to fine的思想，构建粗采样网络和细采样网络可以更好的对空间进行采样。
随机采样：将射线从近场到远场的范围[t_n,t_f][tn,tf]均匀划分成N_cNc个区间，在每段区间内随机取一个点，将其空间坐标\boldsymbol{r}r和空间视角\boldsymbol{d}d输入粗采样网络，得到粗采样网络预测的该空间点的RGB\sigmaRGBσ。
重要性采样：粗采样网络的输出中包括一条射线上所有点的权重w_iwi（具体将在下一节解释），将其归一化后作为采样区间的概率密度函数PDF，按照概率密度函数随机采样N_fNf个点，与前面分段均匀采样的N_cNc个点合并后输入细采样网络。
NeRF将粗采样网络和细采样网络的渲染结果（2D像素点的RGB值），分别与ground truth计算均方误差，将两者之和作为总的loss，来同时训练两个网络。
5. 体绘制原理
得到采样点的密度和RGB值之后，NeRF使用计算机图形学中经典的体绘制技术（Voulume rendering）将3D采样值渲染至2D平面，具体原理如下。
光在介质中的衰减满足以下微分方程：
dL=-L\cdot \sigma \cdot dtdL=−L⋅σ⋅dt
其中LL为光强，\sigmaσ为衰减系数，其解为：
L=L_0\cdot \exp(-\int{\sigma}dt)L=L0⋅exp(−∫σdt)
假设：1) 空间点的颜色\boldsymbol{c}=[R,G,B]^Tc=[R,G,B]T和视线方向\boldsymbol{d}d有关；2) 空间点的光密度\sigmaσ和视线方向\boldsymbol{d}d无关。因为观察到的物体颜色会受到观察视角的影响（比如金属反射面），而光密度是由物体材质所决定。这一点假设也体现在了NeRF的网络结构当中。
根据光在介质中的衰减方程和体绘制原理，对于一个像素，其渲染颜色为
\boldsymbol{C}=\int_{t_{n}}^{t_{f}} T(t) \sigma(\boldsymbol{r}(t)) c(\boldsymbol{r}(t), \boldsymbol{d}) dtC=∫tntfT(t)σ(r(t))c(r(t),d)dt
其中，T(t)=\exp(-\int_{t_n}^{t}{\sigma(\boldsymbol{r}(s))}ds)T(t)=exp(−∫tntσ(r(s))ds)代表[t_n,t][tn,t]内的累积透光率，\sigma(\boldsymbol{r}(t))dtσ(r(t))dt代表距离微元dtdt内的光强衰减率，等效为\boldsymbol{r}(t)r(t)处的反射率。
再对积分式进行离散化处理，使之适于计算机处理：沿着射线取NN个点，在每个点ii所代表的区间长度\delta_iδi内，\sigma, \mathbf{c}σ,c视作常数，得到渲染得到的2D像素RGB值表达式为
\hat{C}(\mathbf{r})=\sum_{i=1}^{N}{T_i(1-\exp(-\sigma_i \delta_i))\mathbf{c_i}}^(r)=i=1∑NTi(1−exp(−σiδi))ci
其中，累计透光率为T_i=\exp(-\sum_{j=1}^{i-1}\sigma_j \delta_j)Ti=exp(−∑j=1i−1σjδj)
用体绘制处理得到的2D像素RGB值，可用于计算loss，也可用于生成最终的2D新视角图像。同时，也可以得到如下两个量：
- \alpha_i=1-\exp(-\sigma_i \delta_i)αi=1−exp(−σiδi)，为第ii个区间的不透明度。因此，T_i=\prod_{j=1}^{i-1}{(1-\alpha_i)}Ti=∏j=1i−1(1−αi)，为前i-1i−1个区间的累积透明度；
- w_i=T_i\cdot \alpha_iwi=Ti⋅αi，为第ii个点的颜色对渲染颜色的贡献度，也就是权重，可用于计算粗采样网络得到的PDF。
6. NeRF的总体流程
综上所述，NeRF实现场景新视角合成的工作流程如下图所示。
- 输入多视角图片（包括像素坐标、像素颜色），以及相机内参、位姿等数据；
- 根据光线步进法产生射线，用随机采样和重要性采样得到空间采样点的坐标；
- 将空间采样点的坐标以及射线的视角输入含有位置编码的NeRF，得到网络对于空间点RGB\sigmaRGBσ值的预测。
- 根据空间点的RGB\sigmaRGBσ值，用体绘制原理，渲染出射线对应的二维像素点的RGBRGB。
- 将预测、渲染得到的像素点的RGBRGB，与ground truth做MSE loss，训练神经网络。
二、代码流程介绍
项目地址：等算子完善后发布
程序实现环境：mindspore1.6.1+CUDA11.1。
整体实验流程：
1. 数据准备
NeRF的训练，除了对同一场景拍摄的多视角照片外，还需要相机的内参和每张照片的位姿。后者是无法通过直接测量得到的，因此需要使用一定的算法来获取。
COLMAP是一款通用的运动恢复结构 (SfM) 和多视图立体 (MVS) 软件，可用于点云重建。我们用COLMAP获取相机的内参和每张照片拍摄时的位姿，并根据3D特征点坐标，估计相机成像的深度范围，从而确定远、近平面，也就是边界。
操作步骤如下：
-
安装COLMAP，官方地址：https://demuc.de/colmap/
-
打开软件。点击File—New project，导入照片文件夹images，并在其所在目录下创建一个
database文件。![import]()
-
点击
Processing---Feature extraction，提取照片中的特征点。 -
点击
Processing---Feature matching，完成特征点匹配。 -
点击
Reconstruction---Start reconstruction，完成稀疏重建。 -
点击
Reconstruction---Dense reconstruction，点击Undistortion，输出稀疏重建的结果，包括相机内参、照片位姿、三维特征点信息的二进制文件。![bins]()
COLMAP生成的数据为二进制文件，需要将数据重新保存成npy格式的文件poses_bounds.npy，便于后续数据的加载。操作步骤如下：
-
运行
img2pose.py文件，运行前需要配置形参为dense文件夹所在的目录。![gen_pose]()
-
运行后，在相应文件夹下生成
pose_bound.npy文件，数据预处理任务完成。pose_bound.npy文件内容：包含相机的内参、每幅图像的位姿（从相机坐标系到世界坐标系的变换矩阵），以及根据照片中的三维特征点信息计算出的每幅照片的成像深度范围。
此外，原始的拍摄照片像素数太多，无法进行计算，需要事先对图片进行下采样，减少计算量，此处采用8倍降采样。
2. 模型搭建
2.1 主干网络模块
首先是主干模型的搭建。模型输入：空间三维坐标在位置编码后的高维向量；模型输出：相应三维坐标点处的 的预测值。模型的具体架构参考原理部分。
这里需要注意模型权重的初始化，它对训练时收敛的快慢影响很大。这里，采用权重和偏置都在(-\sqrt{k},\sqrt{k})(−k
class LinearReLU(nn.Cell):
def __init__(self, in_channels, out_channels):
super().__init__()
self.linear_relu = nn.SequentialCell([
nn.Dense(in_channels, out_channels,
weight_init=Uniform(-math.sqrt(1. / in_channels)),
bias_init=Uniform(-math.sqrt(1. / in_channels))),
nn.ReLU()
])
def construct(self, x):
return self.linear_relu(x)
class NeRF(nn.Cell):
def __init__(self, D=8, W=256, input_ch=3, input_ch_views=3, skips=[4], use_viewdirs=False):
super(NeRF, self).__init__()
self.D = D
self.W = W
self.input_ch = input_ch
self.input_ch_views = input_ch_views
self.skips = skips
self.use_viewdirs = use_viewdirs
self.pts_layers = nn.SequentialCell(
[LinearReLU(input_ch, W)] +
[LinearReLU(W, W) if i not in self.skips else LinearReLU(W + input_ch, W)
for i in range(D - 1)]
)
self.feature_layer = LinearReLU(W, W)
if use_viewdirs:
self.views_layer = LinearReLU(input_ch_views + W, W // 2)
else:
self.output_layer = LinearReLU(W, W // 2)
self.sigma_layer = nn.SequentialCell([
nn.Dense(W, 1,
weight_init=Uniform(-math.sqrt(1. / W)),
bias_init=Uniform(-math.sqrt(1. / W))) if use_viewdirs
else nn.Dense(W // 2, 1,
weight_init=Uniform(-math.sqrt(1. / (W // 2))),
bias_init=Uniform(-math.sqrt(1. / (W // 2))))