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

内容导读

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

但是现实环境中,物体多以 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),则提供了 camerastransformsso3  共计 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 立体结构 ing 本教程将介绍:

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

备注:限于篇幅有限,本文仅展示部分代码,查看完整代码请点击此处

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 图像构建有纹理的立体形状

查看完整教程请点击这里