从Open3D VoxelGrid提取grid_index速度过慢，求优化方案

我尝试将Blender导出的.obj网格进行体素化，并转换为numpy数组保存为.npy或nii.gz格式（模拟CT三维图像）。当前使用Open3D结合CUDA实现快速体素化，但提取grid_index的循环步骤耗时远超体素化过程，相关代码如下：

import open3d as o3d
import numpy as np

mesh = o3d.io.read_triangle_mesh('path/to/file.obj')
voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh(mesh, voxel_size=0.01)
# 耗时瓶颈步骤
voxels = np.array([voxel.grid_index for voxel in voxel_grid.get_voxels()])

grid_size = voxels.max(axis=0) + 1

voxel_array = np.zeros(grid_size, dtype=np.uint16)
voxel_array[voxels[:, 0], voxels[:, 1], voxels[:, 2]] = 1

np.save('path/to/file_open3d.npy', voxel_array)

希望找到优化该循环步骤的方法。

1. 利用Open3D底层数组接口批量提取 #

Open3D的VoxelGrid.get_voxels()返回的Voxel对象可直接转换为numpy结构化数组，无需Python循环即可批量获取grid_index：

import open3d as o3d
import numpy as np

mesh = o3d.io.read_triangle_mesh('path/to/file.obj')
voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh(mesh, voxel_size=0.01)

# 直接提取所有voxel的grid_index数组，跳过Python循环
voxels = np.asarray(voxel_grid.get_voxels())[:, :3].astype(np.int64)

grid_size = voxels.max(axis=0) + 1
voxel_array = np.zeros(grid_size, dtype=np.uint16)
voxel_array[voxels[:, 0], voxels[:, 1], voxels[:, 2]] = 1

np.save('path/to/file_open3d.npy', voxel_array)

此方法直接调用Open3D的底层数据转换逻辑，避免了Python层面的迭代开销，速度会有数量级提升。

2. 用np.fromiter减少内存开销（兼容旧版本） #

如果因Open3D版本限制无法使用上述方法，可使用np.fromiter从迭代器直接生成数组，无需先创建完整列表：

# 替代列表推导式的高效写法
voxels = np.fromiter(
    (coord for voxel in voxel_grid.get_voxels() for coord in voxel.grid_index),
    dtype=np.int64
).reshape(-1, 3)

这种方式比列表推导式更节省内存，因为它直接从迭代器流式生成数组，不会在内存中暂存完整的列表对象。

3. 直接生成体素数组（跳过显式提取） #

若最终目标是生成voxel_array，可借助VoxelGrid的边界信息直接计算，或利用Open3D的CUDA加速接口直接填充体素（部分新版本支持）：

# 获取体素网格的边界与尺寸
min_bound = voxel_grid.get_min_bound()
max_bound = voxel_grid.get_max_bound()
grid_size = ((max_bound - min_bound) / voxel_grid.voxel_size).astype(np.int64) + 1

voxel_array = np.zeros(grid_size, dtype=np.uint16)
# 若Open3D版本支持，可调用CUDA加速的体素填充API，无需手动提取grid_index

部分新版Open3D提供了直接生成体素掩码的接口，可彻底跳过grid_index提取步骤，进一步提升效率。

内容的提问来源于stack exchange，提问作者adamW