基于PyTorch3D的GeoAI实现【ESRI】

Esri 的 AI 原型团队正在以 PyTorch3D API 的一系列 PR 的形式分享一些功能增强功能。这些功能支持 obj 格式的网格的输入/输出 (I/O)，该网格具有多个纹理和代表真实世界几何形状的顶点坐标。对于 GeoAI 任务，这些功能支持跨网格分割管道的任务，例如创建训练数据、提取特征以及在推理过程中将标签应用于网格面。作为开源贡献，我们希望这些功能能够帮助社区推动 3D 网格分割的发展。

推荐：用 NSDT编辑器快速搭建可编程3D场景

我们将功能增强分为 ArcGIS/PyTorch3D 分支中的三个分支，因为它们包含对当前主分支的正交更改。在这三个分支中，我们建议将分支 3 --- multitexture-obj-high- precision 作为主要分支，考虑集成到 PyTorch3D 主分支中，因为它包含分支 1 和 2 中的所有更改以及对浮点 64 中顶点的支持。树枝可以独立站立。

分叉：github.com/Esri/pytorch3d
分支 1 (PR#1572)：multitexture-obj-io-support
分支 2 (PR#1573)：multitexture-obj-point-sampler
分支 3 (PR#1574)：multitexture-obj-high-precision
演示笔记本：github.com/Esri/pytorch3d/.../tutorials/...

1、Esri AI原型团队简介

在 Esri，我作为 AI 原型团队的数据科学家和产品工程师从事研发工作。我的关键任务之一是开发一个数据工程管道，使我的队友能够为各种 GeoAI 工作流程操作网格。例如，在一个工作流程中，我们将点分割模型应用于从网格面采样的点云。在其他情况下，我们直接在网格上进行操作。一般来说，我们的目标是生成一个快速、准确和端到端的工作流程来学习模型，帮助我们将城市景观网格划分为语义类别，例如建筑、树木或我们选择关注的任何其他类别。

2、用PyTorch3D实现GeoAI存在的问题

在许多可用于处理 3D 数据结构的开源库中，哪些是最适合的？尽管这是我们经常评估的一个主题，但我们发现 PyTorch3D 自然而然地适合，因为我们喜欢使用 PyTorch 和其他库（例如 PyTorch Lightning）。毕竟，PyTorch3D 简介通过"PyTorch3D 为使用 PyTorch3D 的 3D 计算机视觉研究提供高效、可重用的组件"说明了您需要了解的所有原因。

确实，PyTorch3D 在 3D 应用深度学习方面帮助我们做出了惊人的事情。然而，尽管有许多很棒的功能，我发现当前的开源库缺乏一些关键功能。例如，PyTorch3D 不完全支持引用多个纹理文件的网格。此外，如果使用 PyTorch3D 从网格中采样点云，则很难将每个点链接到从中采样的网格面。

此外，当使用可以跨越多个城市街区或跨区域多个城市的网格时，处理所有可用纹理文件并将点链接到面的能力是关键要求（仅举几例）。

上面由作者生成并在 ArcGIS Pro 中渲染的图像，代表慕尼黑城市场景的输入网格，该图像将网格中的采样纹理与 Esri 的 PyTorch3D 分支进行比较。给定输入网格（图 1a），我们应该期望采样点云呈现相似的颜色。然而，PyTorch3D 并不完全支持多个纹理（图 1b）。借助 Esri 的 PR 多纹理支持和高精度采样，图 1c 演示了我们如何完整处理网格纹理。

3、Esri AI原型团队的解决方案

首先，我们围绕 PyTorch3D 构建了辅助函数来满足我们处理多个纹理的需求。后来，我们开始使用自定义函数修补 PyTorch3D，以生成将采样点链接到其原始面的张量。然而，我们的团队很快意识到，除非我们实现了可以在 PyTorch3D 内部维持的解决方案，否则自定义补丁无法扩展 - 所以这正是我们最终所做的。

本节提供功能增强的详细摘要。

上图中作者用 Python 生成的图像代表 Hessigheim 城市场景的输入网格，该图像将网格中的采样纹理与 Esri 的 PyTorch3D 分支进行比较。图 2a 中的点云表示建筑物表面植被的不良纹理采样结果，而图 2b 中的点云表示预期结果。

下面是功能增强的高级摘要：

multitexture-obj-io-support：该分支通过修改 pytorch3d.io.obj_io 为 obj 网格建立多纹理支持。目前，PyTorch3D 不完全支持使用多个纹理文件读取和写入 obj 网格。在这种情况下，只有许多纹理中的第一个被读入内存。对于包含不同纹理和多个文件的网格，纹理采样的结果可能会导致不良结果，即植被纹理采样到建筑物表面（图 2a 和 2b）。
具体来说，我们创建了 pytorch3d.io.obj_io.subset_obj 并修改了 pytorch3d.io.obj_io.save_obj 来实现这些功能。此外，还以 pytorch3d.utils.obj_utils 的形式提供了一个新的实用函数，它整合了用于支持 obj 子设置和验证操作的多个辅助函数。

此分支和以下分支已更新，以支持 API 的最新更改，从 PyTorch 版本 0.7.4 开始，现在支持面垂直法线的 I/O。解决了多个现有问题，包括#694、#1017 和#1277。

multitexture-obj-point-sampler：此分支包含 multitexture-obj-io-support 中的所有更改，并添加了对直接从 obj 格式的网格采样点的支持。该分支引入了一个新函数 pytorch3d.ops.sample_points_from_obj，它利用了 pytorch3d.ops.sample_points_from_meshes 中已存在的核心函数。
直接从具有许多大型纹理文件的 obj 中采样点可能比网格数据结构更有优势，因为网格结构在内存中连接纹理。有三个关键特征需要强调。

首先，该分支允许sample_points_from_meshes 和sample_points_from_obj 返回一个映射器张量，该张量将每个采样点链接到采样的面。

其次，该分支允许强制sample_points_from_obj从所有面返回至少一个点，无论面面积如何（sample_all_faces=True）。

第三，该分支允许用户使用 min_sampling_factor 而不是固定数量的点来指定所需点云的密度。链接到问题#1571。

multitexture-obj-high- precision：通过以浮点精度 64（即双张量）读取 obj 顶点来添加对 pytorch3d.io.obj_io 和 pytorch3d.ops.sample_points_from_obj 的支持。
尽管 PyTorch3D 目前允许控制小数位的输出，但如果 obj 中网格的顶点坐标基于真实世界坐标，则顶点值有可能会损失显着的数字精度。实际上，这意味着网格面的顶点坐标可能会偏移一米或更多。链接到问题#1570。

4、增强3D解析能力的示例

本部分介绍从带有和不带有 multitexture-obj-high- precision 高精度分支的 obj 文件读取的网格顶点示例。

在下面的代码片段中，我们提供了从同一 obj 文件读取两个张量的示例。第一个张量显示每个 y 顶点值（中间列舍入到最接近的 0.5 值。第二个张量显示位于或接近其实际位置的相同顶点值，没有舍入。作为参考，源 obj 具有空间 ETRS 1989 UTM Zone 32 的参考（有关更多信息，请参阅developers.arcgis.com 上的空间参考）。

复制代码

# example obj verts sampled without high-precision fork
tensor([
      [692202.1875, 5336173.5000, 512.4358],     
      [692202.8125, 5336174.0000, 512.1414],
      [692202.1250, 5336174.0000, 512.5216],   
      ...,
      [692324.2500, 5336135.5000, 510.2307],
      [692349.1875, 5336142.0000, 526.5537], 
      [692352.4375, 5336169.5000, 522.2985]]
)
     
# example obj verts sampled with high-precision fork
tensor([
      [692202.1617, 5336173.5623, 512.4358], 
      [692202.7912, 5336174.0104, 512.1414],
      [692202.1285, 5336174.0917, 512.5216], 
      ..., 
      [692324.2332, 5336135.5843, 510.2307], 
      [692349.1590, 5336142.1059, 526.5537], 
      [692352.4089, 5336169.6091, 522.2984]],
      dtype=torch.float64
)

功能增强的详细摘要：

pytorch3d.ops.sample_points_from_obj() 是一个新函数，允许用户使用新的自动采样功能从所有面上采样至少一个点，该功能确定要采样的点的数量。

尽管是一个新函数，sample_points_from_obj 重新封装了 pytorch3d.ops.sample_points_from_meshes() 中的现有 PyTorch3D 功能。重要的是，sample_points_from_obj 中的增强功能允许使用最小采样因子和点到面索引映射器张量对所有面进行采样，该张量允许将每个点链接到其原始面。
pytorch3d.ops.sample_points_from_meshes() 被修改为通过将关键功能分组到可在sample_points_from_obj 中利用的辅助函数中来启用sample_points_from_obj()。

此外，sample_points_from_meshes 被稍微修改为可选地返回点到面索引映射器，这可以允许用户恢复每个点的面；但是，对sample_points_from_meshes 的修改保持在最低限度，并且不提供sample_points_from_obj 中的新功能。
pytorch3d.io.obj_io.subset_obj() 是一个新函数，允许用户根据选定的面索引对 obj 网格进行子集化。例如，如果工作流预测每个面的分类，则可以使用此函数仅针对这些面对网格进行子集化。
pytorch3d.io.obj_io.save_obj() 和 pytorch3d.io.obj_io.load_obj_as_meshes() 提供集成的多纹理 obj 支持，允许用户读取和处理所有可用纹理； PyTorch3D 以前仅读取具有多个纹理的输入 obj 文件中的第一个纹理，这可能会导致不良的纹理采样和输出。
pytorch3d.utils.obj_utils 提供了可在 pytorch3d.ops 和 pytorch3d.io.obj_io 中使用的常用实用程序，例如合并 obj 验证 (_validate_obj) 和用于子集 obj 数据的核心实现。
multitexture-obj-high- precision：该分支包含 multitexture-obj-io-support 和 multitexture-point-sampler 的所有功能，并引入了加载 obj 顶点并将 high_ precision 设置为 True 的选项；这意味着支持真实世界坐标和几何图形的完全精度。

其他实用程序函数经过修改，以支持整个代码库中浮点 32 和 64 的 dtype 期望，以包括 Transform3d 和 Meshes 结构类。这是一个值得注意的功能，因为默认浮点 32 可能会导致网格中每个顶点的精度显着损失。由于支持真实世界坐标，我们建议使用此高精度分支而不是其他两个分支。

5、新增API的演示代码

为了快速参考，本部分包含演示 API 的一些关键更改的代码片段。

安装和引入库：

复制代码

# install direct from our forked repository: https://github.com/Esri/pytorch3d/tree/multitexture-obj-high-precision
!pip install 'git+https://github.com/Esri/pytorch3d.git@multitexture-obj-high-precision'

# import the forked pytorch3d library and new functions
from pytorch3d.io import load_obj
from pytorch3d.ops import sample_points_from_obj
from pytorch3d.io import subset_obj, save_obj
使用 high-precision 加载obj模型：

# install direct from our forked repository: https://github.com/Esri/pytorch3d/tree/multitexture-obj-high-precision
!pip install 'git+https://github.com/Esri/pytorch3d.git@multitexture-obj-high-precision'

from pytorch3d.io import load_obj

# the current pytorch3d API for load_obj
obj = load_obj(
    path_to_obj_file,
    load_textures=True,
)

# the high-precision branch of the forked pytorch3d API for load_obj
obj = load_obj(
    path_to_obj_file,
    load_textures=True,
    high_precision=True, # the resulting vertex tensors are returned as floating point 64 tensors
)

从obj模型采样点云：

复制代码

# install direct from our forked repository: https://github.com/Esri/pytorch3d/tree/multitexture-obj-high-precision
!pip install 'git+https://github.com/Esri/pytorch3d.git@multitexture-obj-high-precision'

from pytorch3d.ops import sample_points_from_obj

(
    points, # points sampled proportional to face area
    normals, # point normals based on mesh verts
    textures, # point textures sampled from faces
    mappers # an index to each points origin face
) = sample_points_from_obj(
    verts=obj[0],
    faces=obj[1].verts_idx,
    verts_uvs=obj[2].verts_uvs,
    faces_uvs=obj[1].textures_idx,
    texture_images=obj[2].texture_images,
    materials_idx=obj[1].materials_idx,
    texture_atlas=obj[2].texture_atlas,
    sample_all_faces=True, # optionally force sampler to provide at least one point per face
    min_sampling_factor=100, # control how dense the point cloud is where 0 is least dense and 1000 is very dense
    return_mappers=True, # whether to return the point to face mapping
    return_textures=True, # whether to return textures per point
    return_normals=True # whether to return normals per point
)

计算obj模型的子集：

复制代码

# install direct from our forked repository: https://github.com/Esri/pytorch3d/tree/multitexture-obj-high-precision
!pip install 'git+https://github.com/Esri/pytorch3d.git@multitexture-obj-high-precision'

from pytorch3d.io import subset_obj

faces_to_subset = <a torch tensor of face indices to subset>

# all obj objects are subset to the selected face indices
obj_subset = subset_obj(
    obj=obj,
    faces_to_subset=faces_to_subset,
    device=torch.device("cpu") # specify cpu as the target device to subset and return obj sebset
  )

保存obj模型：

复制代码

# install direct from our forked repository: https://github.com/Esri/pytorch3d/tree/multitexture-obj-high-precision
!pip install 'git+https://github.com/Esri/pytorch3d.git@multitexture-obj-high-precision'

from pytorch3d.io import save_obj

# save_obj now allows for saving objs with multiple textures and control over the texture outputs
save_obj(
    f=obj_filename,
    verts=obj[0],
    faces=obj[1].verts_idx,
    verts_uvs=obj[2].verts_uvs,
    faces_uvs=obj[1].textures_idx,
    texture_images=obj[2].texture_images,
    materials_idx=obj[1].materials_idx,
    image_format='jpeg', # or 'png'
    normals=face_verts_normals,
    faces_normals_idx=faces_normals_idx,
    image_quality=95, # we can now choose the image quality according to PIL docs
    image_name_kwargs=dict(
        material_name_as_file_name=True, # if false the obj file name is enumerated to save each material file
        reuse_material_files=True # if true, an obj with multiple subsets that reference the same materials are not written of the same file already exists
    )
)

6、Esri AI原型团队的后续计划

借助此 PyTorch3D 分支和其他工作流程，我们的团队正在积极开发应用各种深度学习工作流程来分割网格和点云的方法。该分支的一个具体应用是为 Hessigheim (H3D) 基准提供的多类分类任务准备数据 - 我们希望很快测试该基准。

7、结束语

支持具有 obj 格式高精度顶点的多纹理网格是在投影到现实世界空间坐标的场景上应用深度学习的关键功能。对于 Esri 的 AI 原型团队，我们修改了 PyTorch3D 以提供这些功能，并将这些修改开源。

借助我们的 PyTorch3D fork，人们可以将点云分割应用于网格分割问题或操作 obj 数据结构中的多纹理网格和标签。这些是所有感兴趣的用户的关键功能，从长远来看，这些功能可能作为主要 PyTorch3D 库的一部分得到最好的维持。

原文链接：用PyTorch3D实现GeoAI --- BimAnt