PyTorch3D立体隐式形状渲染：教你构建场景3D结构

taotao_2016 2022-01-09

展开全文

内容导读
3D 深度学习一直是机器视觉领域的难点，为了准确高效地建立场景的立体模型，得到相对真实的渲染成果，行业内的一些大厂先后开源了自己的研发成果。

本文首发自微信公众号「PyTorch 开发者社区」

随着计算机视觉领域相关技术的发展，针对 2D 图像的处理效率和精度得到了显著地提升。

但是现实环境中，物体多以 3D 立体结构的形式存在，如何准确地提升 AI 系统对复杂现实环境的感知和理解能力，正确处理 3D 图像，正在成为日趋关键的技术难点。

2019 年 Facebook AI 发布 Mesh R-CNN 模型进行 3D 目标检测与形状预测

2020 年 1 月 23 日，Facebook AI 发布了 PyTorch3D v0.1.0 版本。PyTorch3D 是 PyTorch 中经过优化的、高效、可重用组件库，它具有 3 大突出特点：高效、模块化和可微分，旨在简化在 PyTorch 中进行 3D 深度学习的难度。

PyTorch3D 中提供了 3D 算子和渲染两大组件。

在 3D 算子中，Fit Mesh 可以利用 3D 损失函数，把初始的通用形状变形为目标形状，并借助一些规则让目标形状变得更为流畅。而 3D 算子中的光束平差法（Bundle Adjustment），则提供了 cameras、transforms、so3 共计 3 个 API，根据给定相机的视角，形成对照相机视角的映射，从而推断场景的 3D 结构。

渲染则包括纹理网格渲染器（Render Textured Meshes）、DensePose 网格渲染器（Render DensePose Meshed）、彩色点云渲染器（Render Colored Pointclouds）等，借助这些渲染器，可以进一步优化形成的场景 3D 结构。

2020 年 2 月 6 日，PyTorch3D 相关代码在 GitHub 开源。经过 5 个版本的迭代后，2021 年 2 月 9 日，PyTorch3D 发布第 6 个公开版本 v0.4.0，新增隐函数、立体渲染和 NeRF 重新实现等功能，为 3D 深度学习研究提供了更迅速、更灵活的开源库。

图为 PyTorch3D logo，由 PyTorch3D 的隐式立体渲染器生成的

隐式形状渲染（Implicit Shape Rendering）

隐式形状渲染是指基于输入场景的新视角，生成 3D 场景的真实渲染，其核心思想是利用神经网络与可微分渲染，重建 3D 场景表面的隐式形态，这使得仅依靠 2D 视图就可以学习 3D 场景中的几何形状。

进行隐式形状渲染需要几个关键组件，包括数据卷的抽象类（abstraction for volume data）以及可微分的隐式形状渲染器。

为了让行业从业者更容易地针对隐式形状渲染进行尝试，PyTorch3D 已经为用户提供了一系列常用的 3D 算子（3D operators）和损失函数，以及一个模块化、可微分的渲染 API。在指出核心可重用组件的同时，也提供了这些组件经验证的、标准化实现方法。

在 PyTorch3D v0.4.0 中，包括 5 个支持隐式形状渲染的新特性：

1、新增数据卷结构（Volumes data structure），支持 3D 卷的批处理和坐标框架之间的转换；

2、新增多个光线纹理实现方法：

GridRaysampler

MonteCarloRaysampler

NDCGridRaysampler

3、新增多个 Raymarcher 实现方法：

AbsorptionOnlyRaymarcher

EmissionAbsorptionRaymarcher

4、新增隐式渲染器（ImplicitRenderer）和体积渲染器（VolumeRenderer）API，构成 Raysampler 和 Raymarcher

5、新增多个效用函数，如点云到体积的可微分转换。

利用 PyTorch3D 生成的甜甜圈 3D 图像

要使用这些新组件，可以借助一个模块化、文档完备的 NeRF 重新实现。

NeRF 是一个深度学习模型，由 Google Research 团队开发，旨在借助神经辐射场（Neural Radiance Fields）表示场景，从而进行视图合成（View Synthesis）。
NeRF 仅使用非结构化图像集合，就能合成复杂的 3D 场景图。

而改良版的 NeRF 重新实现，性能得到了极大提升，在保证输出图像质量的同时，比正式版本运行得更快。

使用 PyTorch3D 的 NeRF 重新实现功能，生成了形状和光影复杂的 3D 图像示例

教程（Fit Textured Volume）

我们基于 PyTorch3D GitHub 官方教程 Fit Textured Volume，进行了汉化和整理，演示如何在 PyTorch3D 中，利用可微分立体渲染，依据给定场景的一组视图，预测场景的立体结构。

用 Raymarching 构建场景 3D 立体结构

本教程将介绍：

如何创建一个可微分的立体渲染器；
如何创建一个立体模型（包括如何使用 Volumes 类）；
使用可微分的立体渲染器，根据图像拟合立体结构；
将预测的立体结构可视化。

备注：限于篇幅有限，本文仅展示部分代码，完整代码请移步：

https:///console/openbayes/containers/oAYzp70wwPk/overview

0. 安装和导入模块

确保已安装 torch 和 torchvision 。

如果没有安装 pytorch3d ，请使用以下代码安装。

1. 生成场景及 mask 的图像

以下代码将生成本次训练数据。它会通过 fit_textured_mesh.ipynb 教程，从多个角度渲染奶牛图像，并返回以下内容：

一系列由奶牛网格渲染器生成的图像及其剪影的张量；与所有摄像机镜头一一对应的渲染器。

备注：在 generate_cow_renders 函数中实现网格渲染的工作原理请参考：
fit_textured_mesh.ipynb

target_cameras, target_images, target_silhouettes = generate_cow_renders(num_views=40)print(f'Generated {len(target_images)} images/silhouettes/cameras.')

2. 初始化体积渲染器

初始化体积渲染器会从目标图像的每个像素发出一条射线，并沿射线采样一组间隔均匀的点。与每个射线点相对应的密度值和色值，可通过查询场景的体积模型中的相应位置获得。
渲染器由一个 raymarcher 和一个 raysampler 构成。

raysampler 负责从图像像素中发射射线，并沿着射线对点进行取样。此处使用的是 NDCGridRaysampler ，它符合标准的 PyTorch3D 坐标网格规范。

raymarcher 获得射线采样的密度和颜色，并将所有射线渲染成光线源像素的颜色和不透明度值。此处使用的是
EmissionAbsorptionRaymarcher ，它实现了标准的 Emission-Absorption Raymarching 算法。

# render_size 表示渲染图像各个边的像素大小，将其设置为与目标图像尺寸一致# 也就是说将其渲染成与基准图像一样的尺寸render_size = target_images.shape[1]# 渲染场景以（0,0,0）为中心，被限定在一个边长约等于 3.0 (国际单位)的边框内。volume_extent_world = 3.0# 1) 实例化 raysampler# 此处 NDCGridRaysampler 会生成一矩形图像网格的射线，其坐标遵循 pytorch3d 坐标规定# 由于此处设定的体积是 128^3，因此取样 n_pts_per_ray=150# 大致相当于每个体素都有一个射线点# 进一步设置 min_depth=0.1，因为相机平面内的所有表面都超过了 0.1 单位raysampler = NDCGridRaysampler(    image_width=render_size,    image_height=render_size,    n_pts_per_ray=150,    min_depth=0.1,    max_depth=volume_extent_world,)# 2) 实例化 raymarcher.# 此处用的是标准 EmissionAbsorptionRaymarcher # 它会沿着每条射线前进# 将每条射线都渲染成一个单一的 3D 颜色向量和一个不透明度标量raymarcher = EmissionAbsorptionRaymarcher()# 最后，用 raysampler 和 raymarcher 实例化体积渲染器renderer = VolumeRenderer(    raysampler=raysampler, raymarcher=raymarcher,)

3. 初始化体积模型

接下来实例化场景的体积模型。这会使得 3D 空间量化为体积像素，其中每个体素都用体素 RGB 颜色的 3D 向量，和描述体素不透明度的密度标量（范围在[0-1]之间，数字越大不透明越高）来表示。

为了保证密度和颜色的取值范围在 [0-1] 之间，我们会在对数空间中表示体积颜色和密度。模型运行 forward 函数时， log-space 值会通过 sigmoid 函数传递，从而使得 log-space 值处于正确的取值范围。

此外， VolumeModel 还包含渲染器对象。这个对象在整个优化过程中保持不变。

本部分代码还定义了 huber 损失函数，它可以计算出渲染色和 mask 之间的差异。

4. 体积拟合

这一步，我们用可微分渲染来进行体积拟合。

为了拟合体积，我们从 target_camera 的视角进行渲染，并将渲染结果与观察到的 target_images 和 target_silhouettes 进行对比。

这种对比是通过评估的
target_images/rendered_images 和
target_silhouettes/rendered_silhouettes 之间的平均 huber（smooth-l1）误差来完成的。

5. 将优化后的场景体积进行可视化

最后，旋转场景体积的 y 轴，从多个视点进行渲染，将优化后的体积进行可视化。

def generate_rotating_volume(volume_model, n_frames = 50):    logRs = torch.zeros(n_frames, 3, device=device)    logRs[:, 1] = torch.linspace(0.0, 2.0 * 3.14, n_frames, device=device)    Rs = so3_exponential_map(logRs)    Ts = torch.zeros(n_frames, 3, device=device)    Ts[:, 2] = 2.7    frames = []    print('Generating rotating volume ...')    for R, T in zip(tqdm(Rs), Ts):        camera = FoVPerspectiveCameras(            R=R[None],             T=T[None],             znear = target_cameras.znear[0],            zfar = target_cameras.zfar[0],            aspect_ratio = target_cameras.aspect_ratio[0],            fov = target_cameras.fov[0],            device=device,        )        frames.append(volume_model(camera)[..., :3].clamp(0.0, 1.0))    return torch.cat(frames)    with torch.no_grad():    rotating_volume_frames = generate_rotating_volume(volume_model, n_frames=7*4)image_grid(rotating_volume_frames.clamp(0., 1.).cpu().numpy(), rows=4, cols=7, rgb=True, fill=True)plt.show()