脑补出新视角,一个统一的NeRF代码库框架已开源

假设一个物体你看了几张照片后,能想象出其它角度看上去的感觉吗?人是可以做到的,我们能自行推测出没见过的部分,或者说没见过的角度是什么样的。模型其实也有办法做到,给定一些场景图片,它也能脑补出一些未见过角度的图像。

渲染新视角,近来最引人瞩目的就是 ECCV 2020 最佳论文荣誉提名的 NeRF (Neural Radiance Field)了,它不需要以前复杂的三维重建过程,只需要几张照片及 拍摄该照片时相机的位置,就能合成新视角下的图像。NeRF 惊艳的效果吸引了很多视觉方面的研究者,后续做出了一系列优秀的工作。

但困难的地方在于,这样的模型构建起来比较复杂,目前也没有一个统一的代码库框架来实现它们,这无疑会阻碍该领域的进一步探索与发展。为此,OpenXRLab 渲染生成平台构建出高度模块化的算法库 XRNeRF,帮助快速实现 NeRF 类模型的构建、训练与推理。

脑补出新视角,一个统一的NeRF代码库框架已开源

开源地址:https://github.com/openxrlab/xrnerf

什么是 NeRF 类模型

NeRF 类任务,一般指的是在已知视角下捕获场景信息, 包括拍摄到的图像,以及每张图像对应的内参外参,从而合成新视角下的图像。借助 NeRF 论文中的图,我们能很清晰地理解这种任务。

脑补出新视角,一个统一的NeRF代码库框架已开源

选自 arxiv: 2003.08934。

NeRF 在收集图像时会同时收集 5 维场景信息,即一张图像对应一个三维坐标值及另外两个光线辐射角度。这样的场景会通过多层感知机建模为 Radiance Field,也就是说该多层感知机将输入三维坐标点并映射为该点的 Density 和 RGB 颜色,从而利用体素渲染(Volume Rendering)将 Radiance Field 渲染为照片级的虚拟视角。

如上图所示,通过一些图片构建 Radiance Field 后,就能生成新视角下架子鼓的图像。因为 NeRF 并不需要显式地进行三维重建就能得到想要的新视角,所以它提供了一种基于深度学习的三维隐式表征范式,仅使用 2D 的 posed images 数据就能训练出包含 3D 场景信息的深度神经网络。

自 NeRF 之后,类 NeRF 的模型就层出不穷:Mip-NeRF 利用椎体而不是射线从而优化精细结构的生成效果;KiloNeRF 采用数以千计的微型多层感知机而不是单个大型多层感知机,降低计算量,实现实时渲染能力;此外 AniNeRF 和 Neural Body 等模型从简短视频帧中学习人体视角变换,得到很好的视角合成和驱动效果;此外,GN’R 模型利用稀疏视角图片与几何先验,实现不同 ID 间可泛化的人体渲染。

脑补出新视角,一个统一的NeRF代码库框架已开源

GN’R 提出的可泛化人体隐式场表征,实现的单模型人体渲染效果

为 NeRF 装上轮子

尽管当前 NeRF 类算法在研究领域具有非常高的热度,但是毕竟属于比较新的方法,所以模型实现上肯定是要麻烦一些的。如果是用 PyTorch 或 TensorFlow 这样常规的框架,那就首先得找个相近的 NeRF 模型,再在其基础上修改。

这样做会带来几个明显的问题,首先即我们得完全读懂一份实现,才能在其基础上改成我们想要的样子,这一部分其实工作量还是不小的;其次因为不同论文的官方实现并不统一,对比不同 NeRF 类模型源码时就会消耗比较大的精力,毕竟谁也不知道某篇论文的训练过程中是不是有一些新颖的 Trick;最后如果没有一套统一的代码,验证新模型的新想法无疑会慢很多。

为了解决众多问题,OpenXRLab 为 NeRF 类模型构造出一种统一的、高度模块化的代码库框架 XRNeRF。

脑补出新视角,一个统一的NeRF代码库框架已开源

XRNeRF 实现了众多 NeRF 模型,上手更为容易,可轻松复现相应论文的实验结果。XRNeRF 将这些模型分成 datasets、mlp、network、embedder 和 render 这 5 个模块。XRNeRF 的易用性在于,只需要通过 config 机制即能组装不同模块而构成完整模型,极其简单易用,同时也极大地增加了复用性。

在保证易用性的基础上,同时还需要灵活性,XRNeRF 通过另一套注册器机制,能定制化不同模块的具体特性或实现,从而使 XRNeRF 解耦性更强,代码也更易于理解。

此外,XRNeRF 所有实现的算法都是采用 Pipeline 的模式,数据上的 Pipeline 读取原始数据,经过一系列处理后获得模型的输入,模型的 Pipeline 则对输入的数据进行处理,获得对应的输出。这样的 Pipeline 将 config 机制与注册机机制串起来,组成了一个完整的架构。

XRNeRF 实现了众多核心 NeRF 模型,并通过如上三大机制将它们都串起来,构建出既易用又灵活的高度模块化代码框架。

XRNeRF 的核心特性

XRNeRF 是基于 Pytorch 框架的 NeRF 类 算法库,目前已经复现了 scene 和 body 两个方向的 8 篇经典论文。相比直接建模,XRNeRF 在搭模型效率、成本和灵活性上都有显著提升,而且有完善的使用文档、示例和 Issue 反馈机制,概括来说,XRNeRF 的核心特性大概有以下 5 点。

1. 实现了众多主流和核心的算法

例如开山之作 NeRF,CVPR 2021 Best Paper Candidate (NeuralBody)、ICCV 2021 Best Paper Honorable Mention (Mip-NeRF) 和 Siggraph 2022 Best Paper (Instant NGP)。

脑补出新视角,一个统一的NeRF代码库框架已开源

在实现了这些模型的基础上,XRNeRF 还能保证复现效果和论文中的基本一致。如下图所示,从客观的 PSNR 和 SSIM 指标来看,其能很好地复现原版代码的效果。

脑补出新视角,一个统一的NeRF代码库框架已开源

脑补出新视角,一个统一的NeRF代码库框架已开源

2. 模块化设计

XRNeRF 将整个代码框架进行了模块化设计,最大程度地提升了代码的可复用性,便于研究者对现有代码进行阅读和修改。通过分析现有的 NeRF 类模型方法,XRNeRF 设计的具体模块流程如下图所示:

脑补出新视角,一个统一的NeRF代码库框架已开源

模块化的优势在于,假设我们需要修改数据格式,那只需要修改 Dataset 模块下的逻辑,假设我们需要修改渲染图像的逻辑,那就只需要修改 Render model 模块。

3. 标准数据处理管线

针对 NeRF 类算法数据预处理较为复杂和多样的问题,XRNeRF 提供了一套标准数据处理流程。其由多个数据处理操作串行得到,仅需要修改 config 配置文件中的 data pipeline 部分,即可完成数据处理流畅搭建。

脑补出新视角,一个统一的NeRF代码库框架已开源

NeRF 配置数据流程部分。

XRNeRF 中已经实现了多个数据集所需要的数据处理 op,只需要将这些 op 在 config 中按照顺序定义好,即可完成数据处理流程的搭建。如果后续有新的 op 需要加入,也只需要在对应文件夹中完成新 op 的实现,即可一行代码加入到整个数据处理流程当中。

4. 模块化网络构建方式

XRNeRF 中的模型主要由 embedder、MLP 和 render model 组成,并通过 network 连接,这几者之间可以互相解耦,由此可以实现不同算法之间不同模块的替换。

其中 embedder 输入点的位置和视角,输出 embedded 特征数据;MLP 使用 embedder 的输出作为输入,输出采样点的 Density 和 RGB 颜色;render model 则输入 MLP 的输出结果,沿着射线上的点进行积分等操作,从而获得图像上一个像素点的 RGB 值。这三大模块再通过标准的 network 模块连接就构成了完整的模型。

脑补出新视角,一个统一的NeRF代码库框架已开源

自定义 network 模块的代码结构。

5. 良好的复现效果

支持最快 60 秒训练网络,30 帧每秒实时渲染,支持高清晰度、抗锯齿、多尺度场景及人体图像渲染。无论是从客观的 PSNR 和 SSIM 指标还是主观的 demo 展示效果来看,XRNeRF 都能很好地复现原版代码的效果。

XRNeRF 的使用

XRNeRF 框架看起来有非常好的特性,其使用起来也很简单便捷。比如说安装过程,XRNeRF 依赖的开发环境还是比较多的,PyTorch、CUDA 环境、视觉方面的处理库等等。但是 XRNeRF 提供了 Docker 环境,通过 DockerFile 能直接构建镜像文件。

脑补出新视角,一个统一的NeRF代码库框架已开源

我们试了一下,相比一步步配置各种运行环境与包,只需 docker build 一行命令的配置方式显然要方便太多了。此外构建 Docker 镜像时,DockerFile 里面配置了国内镜像地址,所以速度还是很快的,基本不用担心网络问题。

在构建完镜像,并从该镜像启动容器后,我们就能将项目代码,以及数据都通过 docker cp 命令传到容器内。不过也可以直接在创建容器时通过 -v 参数直接将项目地址映射到容器内部。不过这里需要注意的是,数据集是需要放到确定位置的(否则需要改 config 文件),例如 XRNeRF 项目下的 data 文件夹。

一般而言,下载完数据后,大概文件夹结构如下图所示:

脑补出新视角,一个统一的NeRF代码库框架已开源

现在,环境、数据与代码三者都准备好了,只需简短的一行代码,就能执行 NeFR 模型的训练与验证:

python run_nerf.py --config configs/nerf/nerf_blender_base01.py --dataname lego

其中 dataname 表示数据目录下的具体数据集,config 表示模型的具体配置文件。因为 XRNeRF 采用的是高度模块化的设计,其 config 使用字典来构建,虽然乍一看可能会觉得有一点点繁琐,但实际理解了 XRNeRF 的设计结构之后,阅读起来就非常简单了。

主观看上去,config 配置文件(nerf_blender_base01.py)包含了训练模型所有必要的信息,优化器、分布式策略、模型架构、数据预处理与迭代等等,甚至很多图像处理相关的配置也都包含在内。总结来说,除了具体的代码实现,config 配置文件描述了整个训练、推理过程。

脑补出新视角,一个统一的NeRF代码库框架已开源

描述模型结构部分的 Config 配置。

总的体验起来,XRNeRF 从基础的运行环境搭建到最终执行训练任务都是比较流畅的。况且通过配置 config 文件,或者实现具体的 OP,同样也能获得非常高的建模灵活性。相比直接使用深度学习框架建模,XRNeRF 无疑将减少大量的开发工作,研究者或者算法工程师也能花更多时间在模型或任务创新上。

NeRF 类模型目前仍然是计算机视觉领域的研究重点,XRNeRF 这样统一的代码库,就像 HuggingFace 的 Transformer 库一样能聚集越来越多的优秀研究工作,聚集越来越多的新代码与新想法。反过来 XRNeRF 同样也将极大地加快研究者对 NeRF 类模型探索的脚步,便于将这一新领域应用到新场景与新任务中,NeRF 的潜力也将由此加速展开。

© 版权声明

相关文章