[论文解析]OmniRe: Omni Urban Scene Reconstruction

OmniRe: Omni Urban Scene Reconstruction

论文地址：https://arxiv.org/abs/2408.16760

代码地址：https://github.com/ziyc/drivestudio

项目地址：https://ziyc.github.io/omnire/

---

论文解读

总结

1. 这篇论文代表了一种重建的方向，就是对场景内的物体进行细致分类，通过单独的优化提高重建效果，这对于动态场景重建是一条正确的道路，因为在动态场景尤其是自动驾驶的场景中，很多动态物体或者静态背景都是一闪而过的，提高次数但是不区分类别对于最后的效果没有正向收益的。而且通过这种方法可以积累动态或者静态资产，有利于资源重复利用，对于可编辑场景仿真有较大收益
2. 神经场景图，对所有的动态物体进行了显式区分表示，为每个障碍物单独建模并选择合适的表示方式，然后结合背景以及天空节点进行重建；
3. 车辆使用静态高斯重建，并使用刚体变换表示随时间的运动，而且对于所有的刚体节点，优化他们每一帧的位姿；
4. 用SMPL模型来对人体进行重建，通过人体关节和姿态进行建模，相当于在训练过程中增加了更多的优化参数。但是在整体前景部分没有考虑光影变化，没有引入shadow的模型，所以在项目页面可以看到，人跳舞的场景是没有光影变化的，不过这个也是比较困难的点，对于自动驾驶场景来说人体光影不是高优任务；
5. 在代码实现上面，这是一篇全面的工作，仓库中集成了几篇非常优秀的工作的手动实现，包括DeformableGS，PVG等，兼容性和扩展性都比较强；
   1. 渲染器采用了gsplat库，没有用原版的diff-gaussian-rasterization库，gsplat对于很多GS方向的新工作支持和更新都是比较快的，很多开源方案都会给gsplat提PR，这个对于特性对比试验是很有好处的；
   2. 当然代码中也存在一些小问题，这些等到读者自己运用的时候再探索吧；

数据集初始化

1. 根据不同数据集读取数据(Waymo, NuScenes, ArgoVerse, PandaSet, KITTI, NuPlan)，正如之前提到的代码写的条分缕析，因此对各个数据集的处理也比较详细，适配自己的数据集也比较方便；
2. 以NuScenes为例，整体读取流程如下图，如果你要适配自己的数据集，我建议先仿照preprocess部分，将你的数据处理为dataloader的读取形式，这样是最快的方法，修改dataloader不是good idea。

训练过程

1. BasicTrainer的初始化，场景初始化，计算场景原点和半径：self._init_scene(scene_aabb=scene_aabb)

2. loss计算函数初始化self._init_losses()

   1. depth的loss有两类三种，通过标志位控制，包括正常深度和逆深度两类，smooth_l1_loss，l1_loss以及mse_loss以及三种loss方案
   2. 天空loss计算为二维交叉熵损失；这里作者实现了一个更安全的BCE算法，规避了Nan值以及梯度异常带来的训练停滞的问题；

3. 图像指标评测计算（SSIM，PSNR，LPIPS）

4. **init_gaussians_from_dataset()**每个节点相当于一个单独的GaussianModel，内部实现为GaussianSplating的表现形式，颜色通过球谐函数表示；

   1. 类的继承关系：VanillaGaussians(nn.Module) > RigidNodes > DeformableNodes / SMPLNodes
   2. 背景初始化为VanillaGaussians，也就是最基础的三维高斯模型，其他模型是在此基础上加入了对应的优化参数生成的；

5. 优化器初始化：initialize_optimizer()

   1. Adam优化器
   2. 混合精度训练torch.cuda.amp.GradScaler
   3. 每个Node设置学习率变化函数，学习率更新函数如下，ramp可配置为sin函数，

   $$lr=\{\begin{array}{ll}\text{lr\_pre\_warmup}+(\text{lr\_init}-\text{lr\_pre\_warmup})\cdot ramp\left(\frac{step}{\text{warmup\_steps}}\right),& \text{if }step<\text{warmup\_steps}\\ \exp \left(\log (\text{lr\_init})\cdot (1-t)+\log (\text{lr\_final})\cdot t\right),& \text{if }step\ge \text{warmup\_steps}\end{array}$$

6. 训练前处理preprocess_per_train_step()，zero grad设置，step设置

7. 训练函数forward()

   1. 根据时间计算帧数，作为rigidNode和deformableNode的输入，这两个节点的优化和时间戳相关
   2. 相机处理process_camera() ，调用CameraOptModule对相机位姿进行优化，embed和decoder的组合优化相机的delta pose
   3. gaussian生成collect_gaussians()
   4. 高斯渲染**render_gaussians()** ：gsplat执行渲染，生成rgb,depth,opacity图像，用于loss计算
   5. affine_transformation来自于**AffineTransform**类，帧数通过embed和decoder输出仿射矩阵，对每一帧的颜色进行修正，避免不同相机曝光不一致导致的效果变差问题；类似于ScaffoldGS中的appearance embedding；

8. 可见性更新update_visibility_filter(self) ：根据node分类，更新当前模型的半径

9. train_step_camera_downscale = trainer._get_downscale_factor() 添加了变分辨率训练，根据步数会下采样图像；

10. loss计算compute_losses()

    1. L1，SSIM的loss计算
    2. sky节点的二维交叉熵损失
    3. depth的loss计算，针对lidar的深度加入了又步长计算的延迟因子；
    4. 对不透明度熵损失加入正则化，采用负对数熵正则化
    5. 逆深度平滑正则化，使得重建的深度更为合理；针对深度图的平滑性损失计算，使用了kornia.losses.inverse_depth_smoothness_loss
    6. 对模型中的仿射变换参数进行约束，使其接近于单位矩阵；
    7. 动态区域范围的L1Loss，在这里动态区域是根据不透明度计算的mask；
    8. RigidNodes正则化
    9. DeformableNodes正则化
    10. VanillaGaussians正则化
    11. SMPLNodes正则化

11. backward()梯度回传，自动求导；获取混合精度训练的梯度缩放比例，如果梯度缩放比例没有被降低，意味着上一次的反向传播是稳定的，因此可以执行学习率调度器的更新。

12. postprocess_per_train_step()后处理，这里相当于3DGS的densify_and_prune() 函数，具体实现在vanilla.py

    1. alpha重置间隔reset_alpha_interval=3000，致密化截止次数15000，prune间隔100；densify_and_prune执行条件：在重置间隔内且大于prune间隔，小于截止iteration；每100次执行一次
    2. split条件（1分2）：平均梯度范数大于densify_grad_thresh=0.0005；高斯点的尺度大于阈值densify_size_thresh=0.002；如果步数小于 stop_screen_size_at=5000，且高斯点的2D半径大于split_screen_size=0.05 ；
    3. clone条件：平均梯度范数大于densify_grad_thresh=0.0005；高斯点的尺度小于等于阈值densify_size_thresh；
    4. cull条件：透明度小于cull_alpha_thresh=0.005 ；scale过大的超过场景尺度相乘系数cull_scale_thresh=0.5 ；2D半径大于阈值cull_screen_size: 0.15 ；（3个条件是or关系）
    5. 每3000次重置不透明度；

一些问题

1. affine_transformation部分是存在问题的，此处的问题可以参考官方仓库的一个issue
2. 这个方案由于不是原版3DGS代码修改的，所以无法导出SIBR Viewer所需要的文件，如果有这方面需求可以参考这个PR，自己修改代码即可；笔者没有尝试nerfstudio的可视化方案，不过估计是差不多的自己动手即可：
3. 代码存在同一个数据重复两次效果差异较大的问题，暂时没找到bug点，所以对于复现可能有些阻碍；