3D模型处理的并行化

今天我们将讨论如何使用 Python 多进程来处理大量3D数据。我将讲述一些可能在手册中找到的一般信息，并分享我发现的一些小技巧，例如将 tqdm 与多处理 imap 结合使用以及并行处理存档。

那么我们为什么要诉诸并行计算呢？使用数据有时会出现与大数据相关的问题。每次我们遇到 RAM 不适合的数据时，我们都需要逐段处理它。幸运的是，现代编程语言允许我们生成在多核处理器上完美工作的多个进程（甚至线程）。注意：这并不意味着单核处理器无法处理多处理，这是有关该主题的 Stack Overflow 讨论。

今天我们将尝试计算网格和点云之间的距离这一常见的 3D 计算机视觉任务。例如，当你需要在所有可用网格中查找定义与给定点云相同的 3D 对象的网格时，可能会遇到此问题。

我们的数据由存储在 .7z 存档中的 .obj 文件组成，这在存储效率方面非常出色。但是当我们需要访问它的确切部分时，我们应该付出努力。在这里，我定义了包装 7-zip 存档并提供底层数据接口的类。

from io import BytesIO
import py7zlib

class MeshesArchive(object):
    def __init__(self, archive_path):
        fp = open(archive_path, 'rb')
        self.archive = py7zlib.Archive7z(fp)
        self.archive_path = archive_path
        self.names_list = self.archive.getnames()
        
    def __len__(self):
        return len(self.names_list)
    
    def get(self, name):
        bytes_io = BytesIO(self.archive.getmember(name).read())
        return bytes_io

    def __getitem__(self, idx):
        return self.get(self.names[idx])
    
    def __iter__(self):
        for name in self.names_list:
            yield self.get(name)

这个类几乎不依赖 py7zlib 包，它允许我们在每次调用 get 方法时解压缩数据，并为我们提供存档内的文件数量。我们还定义了 __iter__ ，它将帮助我们像在可迭代对象上一样在该对象上启动多处理映射。

这个定义为我们提供了迭代存档的可能性，但它是否允许我们并行随机访问内容？这是一个有趣的问题，我在网上没有找到答案，但如果深入研究 py7zlib 的源代码，我们可以回答它。

在这里，我提供了 pylzma 的代码片段：

class Archive7z(Base):
  def __init__(self, file, password=None):
    # ...
    self.files = {}
    # ...
    for info in files.files:
      # create an instance of ArchiveFile that knows location on disk
      file = ArchiveFile(info, pos, src_pos, folder, self, maxsize=maxsize)
      # ...
      self.files.append(file)
    # ...
    self.files_map.update([(x.filename, x) for x in self.files])
        
  # method that returns an ArchiveFile from files_map dictionary
  def getmember(self, name):
      if isinstance(name, (int, long)):
          try:
              return self.files[name]
          except IndexError:
              return None

      return self.files_map.get(name, None)
    
    
class Archive7z(Base):
  def read(self):
    # ...
    for level, coder in enumerate(self._folder.coders):
      # ...
      # get the decoder and decode the underlying data
      data = getattr(self, decoder)(coder, data, level, num_coders)

    return data

摘自pylzma源码，省略了很多

我相信从上面的要点可以清楚地看出，只要同时多次读取存档，就没有理由被阻止。

接下来我们快速介绍一下什么是网格和点云。首先是网格，它们是顶点、边和面的集合。顶点由空间中的 (x,y,z) 坐标定义，并分配有唯一的编号。边和面相应地是点对和三元组的组，并使用提到的唯一点 ID 进行定义。通常，当我们谈论“网格”时，我们指的是“三角形网格”，即由三角形组成的表面。使用 trimesh 库在 Python 中处理网格要容易得多，例如它提供了在内存中加载 .obj 文件的接口。要在 Jupyter Notebook 中显示 3D 对象并与之交互，可以使用 k3d 库。

因此，通过以下代码片段，我回答了这个问题：“如何使用 k3d 在 jupyter 中绘制 atrimeshobject？”

import trimesh
import k3d

with open("./data/meshes/stanford-bunny.obj") as f:
    bunny_mesh = trimesh.load(f, 'obj')

plot = k3d.plot()
mesh = k3d.mesh(bunny_mesh.vertices, bunny_mesh.faces)
plot += mesh
plot.display()

k3d 显示的斯坦福兔子网格（不幸的是这里没有响应）

其次，点云，它们是表示空间中物体的 3D 点阵列。许多 3D 扫描仪生成点云作为扫描对象的表示。为了演示目的，我们可以读取相同的网格并将其顶点显示为点云。

import trimesh
import k3d

with open("./data/meshes/stanford-bunny.obj") as f:
    bunny_mesh = trimesh.load(f, 'obj')
    
plot = k3d.plot()
cloud = k3d.points(bunny_mesh.vertices, point_size=0.0001, shader="flat")
plot += cloud
plot.display()

将顶点绘制为点云

k3d绘制的点云

正如上面提到的，3D 扫描仪为我们提供了点云。假设我们有一个网格数据库，并且希望在数据库中找到与扫描对象（即点云）对齐的网格。为了解决这个问题，我们可以提出一种简单的方法。我们将搜索给定点云的点与存档中的每个网格之间的最大距离。如果对于某些网格来说，1e-4 的距离较小，我们会认为该网格与点云对齐。

最后，我们来到了多处理部分。请记住，我们的存档有大量文件可能无法同时放入内存中，我们更喜欢并行处理它们。为了实现这一点，我们将使用多处理池，它使用 map 或 imap/imap_unordered 方法处理用户定义函数的多次调用。 map 和 imap 之间影响我们的区别在于， map 在发送到工作进程之前将可迭代对象转换为列表。如果存档太大而无法写入 RAM，则不应将其解压到 Python 列表中。在另一种情况下，它们的执行速度相似。

[Loading meshes: pool.map w/o manager] Pool of 4 processes elapsed time: 37.213207403818764 sec
[Loading meshes: pool.imap_unordered w/o manager] Pool of 4 processes elapsed time: 37.219303369522095 sec

在上面你可以看到从适合内存的网格存档中进行简单读取的结果。

使用 imap 更进一步。让我们讨论如何实现找到靠近点云的网格的目标。这是数据，我们有来自斯坦福模型的 5 个不同的网格。我们将通过向斯坦福兔子网格的顶点添加噪声来模拟 3D 扫描。

import numpy as np
from numpy.random import default_rng

def normalize_pc(points):
    points = points - points.mean(axis=0)[None, :]
    dists = np.linalg.norm(points, axis=1)
    scaled_points = points / dists.max()
    return scaled_points


def load_bunny_pc(bunny_path):
    STD = 1e-3 
    with open(bunny_path) as f:
        bunny_mesh = load_mesh(f)
    # normalize point cloud 
    scaled_bunny = normalize_pc(bunny_mesh.vertices)
    # add some noise to point cloud
    rng = default_rng()
    noise = rng.normal(0.0, STD, scaled_bunny.shape)
    distorted_bunny = scaled_bunny + noise
    return distorted_bunny

当然，我们之前对下面的点云和网格顶点进行了标准化，以在 3D 立方体中缩放它们。

为了计算点云和网格之间的距离，我们将使用 igl。为了最终确定，我们需要编写一个将在每个进程及其依赖项中调用的函数。让我们用下面的片段来总结一下。

import itertools
import time

import numpy as np
from numpy.random import default_rng

import trimesh
import igl
from tqdm import tqdm

from multiprocessing import Pool

def load_mesh(obj_file):
    mesh = trimesh.load(obj_file, 'obj')
    return mesh

def get_max_dist(base_mesh, point_cloud):
    distance_sq, mesh_face_indexes, _ = igl.point_mesh_squared_distance(
        point_cloud,
        base_mesh.vertices,
        base_mesh.faces
    )
    return distance_sq.max()

def load_mesh_get_distance(args):
    obj_file, point_cloud = args[0], args[1]
    mesh = load_mesh(obj_file)
    mesh.vertices = normalize_pc(mesh.vertices)
    max_dist = get_max_dist(mesh, point_cloud)
    return max_dist

def read_meshes_get_distances_pool_imap(archive_path, point_cloud, num_proc, num_iterations):
    # do the meshes processing within a pool
    elapsed_time = []
    for _ in range(num_iterations):
        archive = MeshesArchive(archive_path)
        pool = Pool(num_proc)
        start = time.time()
        result = list(tqdm(pool.imap(
            load_mesh_get_distance,
            zip(archive, itertools.repeat(point_cloud)),
        ), total=len(archive)))
        pool.close()
        pool.join()
        end = time.time()
        elapsed_time.append(end - start)

    print(f'[Process meshes: pool.imap] Pool of {num_proc} processes elapsed time: {np.array(elapsed_time).mean()} sec')
    
    for name, dist in zip(archive.names_list, result):
        print(f"{name} {dist}")
    
    return result
  
 if __name__ == "__main__":
    bunny_path = "./data/meshes/stanford-bunny.obj"
    archive_path = "./data/meshes.7z"
    num_proc = 4
    num_iterations = 3

    point_cloud = load_bunny_pc(bunny_path)
    read_meshes_get_distances_pool_no_manager_imap(archive_path, point_cloud, num_proc, num_iterations)

这里 read_meshes_get_distances_pool_imap 是一个核心函数，其中完成了以下操作：

MeshesArchive 和 multiprocessing.Pool 已初始化
应用 tqdm 来监视池进度，并手动完成整个池的分析
执行结果的输出

请注意我们如何将参数传递给 imap，使用 zip(archive, itertools.repeat(point_cloud)) 从 archive 和 point_cloud 创建新的可迭代对象。这使我们能够将点云数组粘贴到存档的每个条目，从而避免将存档转换为列表。

执行结果如下所示：

100%|####################################################################| 5/5 [00:00<00:00,  5.14it/s]
100%|####################################################################| 5/5 [00:00<00:00,  5.08it/s]
100%|####################################################################| 5/5 [00:00<00:00,  5.18it/s]
[Process meshes: pool.imap w/o manager] Pool of 4 processes elapsed time: 1.0080536206563313 sec
armadillo.obj 0.16176825266293382
beast.obj 0.28608649819198073
cow.obj 0.41653845909820164
spot.obj 0.22739556571296735
stanford-bunny.obj 2.3699851136074263e-05

我们可以发现斯坦福兔子是最接近给定点云的网格。还可以看出，我们没有使用大量数据，但我们已经证明，即使存档中有大量网格，该解决方案也能发挥作用。

多重处理使数据科学家不仅在 3D 计算机视觉方面而且在机器学习的其他领域都取得了出色的表现。理解并行执行比循环内执行要快得多，这一点非常重要。尤其是当算法编写正确时，差异变得非常显着。大量数据揭示的问题如果没有创造性的方法来利用有限的资源就无法解决。幸运的是，Python 语言及其丰富的库可以帮助我们数据科学家解决此类问题。