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

内容导读

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://openbayes.com/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()

6. 结论

本教程中演示了如何优化场景的 3D 立体构造，使已知视点的体积渲染与每个视点的观测图像相匹配。

教程中的渲染是使用 NDCGridRaysampler 和 EmissionAbsorptionRaymarcher 构成的 PyTorch3D 立体渲染器完成的。

为 2D 图像构建有纹理的立体形状

完整教程请查看： https://openbayes.com/console/openbayes/containers/oAYzp70wwPk/overview

参考：

https://ai.facebook.com/blog/-introducing-pytorch3d-an-open-source-library-for-3d-deep-learning/

https://github.com/facebookresearch/pytorch3d