如何用Numpy优化3D点云到2D画布的栅格化效率？

我有一个包含约45000个点的3D点列表，示例如下：

import numpy as np
ps = np.array([[0.77258706, 1.1149821 , 1.0494938 ],
       [0.77258706, 1.1149821 , 1.0494938 ],
       [0.77258706, 1.1149821 , 1.0494938 ],
       [1.3196146 , 1.2129393 , 0.96248001],
       [1.3197775 , 1.2117194 , 0.96213955],
       [1.3199375 , 1.2104989 , 0.96179962],
       [1.3200942 , 1.2092786 , 0.96145791],
       [1.3202487 , 1.2080572 , 0.96111906],
       [1.3204001 , 1.2068353 , 0.96078068]])

需求是将这些点投影到2D画布上：提取x、y坐标缩放后作为画布位置，z分量缩放后设置颜色，所有值需缩放适配画布。现有实现代码如下：

canv = np.zeros(shape=[3840, 2160, 3])
for p in ps:
    x_pos = p[0] * scale_factor
    y_pos = p[1] * scale_factor
    val =   p[2] * color_scale_factor
    canv[int(x_pos)][int(y_pos)][0:3] = colors[int(val)].get_rgb()

当前每帧处理耗时约120ms（帧率8FPS），希望通过Numpy技巧大幅提升处理效率。

优化方案 #

核心思路：替换循环为Numpy向量化操作 #

Python循环处理大量数据时效率极低，Numpy的向量化运算基于底层C实现，能批量完成计算，效率提升量级可达数十倍。具体步骤如下：

1. 批量计算坐标与颜色索引
   直接对整个点云数组做元素级运算，避免逐个点循环：

# 批量计算缩放后的坐标并转为整数索引
x_indices = (ps[:, 0] * scale_factor).astype(np.int32)
y_indices = (ps[:, 1] * scale_factor).astype(np.int32)
# 批量计算颜色索引
color_indices = (ps[:, 2] * color_scale_factor).astype(np.int32)

2. 过滤超出画布范围的点
   避免索引越界错误，同时减少无效运算：

# 生成有效点掩码：确保x、y坐标在画布尺寸范围内
valid_mask = (x_indices >= 0) & (x_indices < canv.shape[0]) & \
             (y_indices >= 0) & (y_indices < canv.shape[1])
# 筛选有效索引
valid_x = x_indices[valid_mask]
valid_y = y_indices[valid_mask]
valid_color_idx = color_indices[valid_mask]

3. 批量映射颜色并赋值
   先将颜色列表预转为Numpy数组，再通过索引批量获取颜色，最后一次性赋值到画布：

# 预转换颜色列表为Numpy数组（假设colors是带get_rgb()方法的颜色对象列表）
color_array = np.array([c.get_rgb() for c in colors], dtype=np.float32)
# 批量获取有效点对应的颜色
point_colors = color_array[valid_color_idx]
# 一次性赋值到画布（保持原代码的索引顺序：x对应画布第一维度，y对应第二维度）
canv[valid_x, valid_y] = point_colors

完整优化代码 #

import numpy as np

# 假设scale_factor、color_scale_factor、colors已提前定义
canv = np.zeros(shape=[3840, 2160, 3], dtype=np.float32)

# 1. 批量计算所有索引
x_indices = (ps[:, 0] * scale_factor).astype(np.int32)
y_indices = (ps[:, 1] * scale_factor).astype(np.int32)
color_indices = (ps[:, 2] * color_scale_factor).astype(np.int32)

# 2. 过滤有效点
valid_mask = (x_indices >= 0) & (x_indices < canv.shape[0]) & \
             (y_indices >= 0) & (y_indices < canv.shape[1])
valid_x = x_indices[valid_mask]
valid_y = y_indices[valid_mask]
valid_color_idx = color_indices[valid_mask]

# 3. 预转换颜色数组并批量赋值
color_array = np.array([c.get_rgb() for c in colors], dtype=np.float32)
canv[valid_x, valid_y] = color_array[valid_color_idx]

额外优化建议 #

• 预生成颜色数组：如果colors列表固定，在程序初始化阶段就转为Numpy数组，避免每次处理都重复生成。
• 数据类型优化：画布和颜色数组使用np.float32而非默认的float64，减少内存占用和运算时间。
• 处理重叠点：若多个点映射到同一像素，原代码会直接覆盖；如果需要叠加效果（如取平均或最大值），可使用np.add.at：

  # 示例：叠加颜色（需确保颜色值在0-1范围内）
  np.add.at(canv, (valid_x, valid_y), point_colors)
  # 归一化到0-1范围
  canv = np.clip(canv, 0, 1)
  

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