如何加速基于强度的点云着色Numpy循环？

问题描述 #

我需要基于强度为点云着色，目前用for循环逐个处理点的强度值生成颜色，代码如下：

import numpy as np
points = np.random.random([128*1200,4])
points_colors = np.zeros([points.shape[0], 3])

for idx, p_c in enumerate(points[:, 3]):
    points_colors[idx, :] = color_map(p_c)
points_colors /= 255.0

使用的基础颜色映射函数：

def color_map( value, minimum=0, maximum=255):
    minimum, maximum = float(minimum), float(maximum)
    ratio = 2 * (value-minimum) / (maximum - minimum)
    b = int(max(0, 255*(1 - ratio)))
    r = int(max(0, 255*(ratio - 1)))
    g = 255 - b - r
    return r, g, b

实际使用的更复杂彩虹色映射函数：

def rainbow_color_map(
    val,
    minval = 0,
    maxval=256,
    normalize=False,
    colors=[(1, 1, 255), (1, 255, 1), (255, 1, 1)] * 10,
):
    i_f = float(val - minval) / float(maxval - minval) * (len(colors) - 1)

    i, f = int(i_f // 1), i_f % 1  # Split into whole & fractional parts.

    (r1, g1, b1), (r2, g2, b2) = colors[i], colors[i + 1]
    if normalize:
        return (
            (r1 + f * (r2 - r1)) / maxval,
            (g1 + f * (g2 - g1)) / maxval,
            (b1 + f * (b2 - b1)) / maxval,
        )
    else: 
        return r1 + f * (r2 - r1), g1 + f * (g2 - g1), b1 + f * (b2 - b1)

当前for循环着色耗时远高于Open3D原生的基于XYZ姿态的着色，希望找到加速基于强度着色的方法，也接受无需将XYZI转成XYZRGB的方案。

加速方案 #

1. 向量化改造颜色映射函数 #

Python逐元素for循环效率极低，核心优化思路是用NumPy向量化操作替代循环，让运算在底层C实现中批量完成。

基础color_map的向量化版本

def vectorized_color_map(values, minimum=0, maximum=255):
    minimum, maximum = float(minimum), float(maximum)
    ratio = 2 * (values - minimum) / (maximum - minimum)
    # 批量计算通道值，用np.maximum替代单元素max
    b = np.maximum(0, 255 * (1 - ratio)).astype(np.int32)
    r = np.maximum(0, 255 * (ratio - 1)).astype(np.int32)
    g = 255 - b - r
    # 按RGB通道拼接成(N,3)数组
    return np.stack([r, g, b], axis=1)

调用方式：

points_colors = vectorized_color_map(points[:, 3]) / 255.0

rainbow_color_map的向量化版本

针对分段插值逻辑，用NumPy批量处理索引和插值计算：

def vectorized_rainbow_color_map(
    values,
    minval=0,
    maxval=256,
    normalize=False,
    colors=[(1, 1, 255), (1, 255, 1), (255, 1, 1)] * 10,
):
    colors = np.array(colors)
    n_colors = len(colors) - 1
    # 批量计算插值比例
    i_f = (values - minval) / (maxval - minval) * n_colors
    # 拆分整数索引和小数插值系数
    i = np.floor(i_f).astype(np.int32)
    # 防止索引越界（处理等于maxval的边界情况）
    i = np.clip(i, 0, n_colors - 1)
    f = i_f - i
    
    # 批量获取前后颜色点并计算插值
    c1 = colors[i]
    c2 = colors[i + 1]
    rgb = c1 + f[:, None] * (c2 - c1)
    
    if normalize:
        rgb = rgb / maxval
    return rgb

调用方式：

points_colors = vectorized_rainbow_color_map(points[:, 3])
# 非normalize模式下按需归一化到0-1区间
if not normalize:
    points_colors /= 255.0

2. 利用Open3D直接绑定强度到颜色通道（无需转XYZRGB） #

Open3D支持直接将点云的强度字段作为颜色来源，不需要手动生成RGB数组，这是最轻量化的方案：

import open3d as o3d

# 创建XYZI点云
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(points[:, :3])

# 直接将强度值重复三次作为灰度颜色
pcd.colors = o3d.utility.Vector3dVector(points[:, 3:4].repeat(3, axis=1))

# 若需要彩色渐变，先对强度做向量化映射再赋值
# mapped_colors = vectorized_rainbow_color_map(points[:,3])
# pcd.colors = o3d.utility.Vector3dVector(mapped_colors)

Open3D内部用底层高效实现处理渲染，完全规避Python层循环开销。

3. 用Numba JIT编译加速循环 #

如果映射逻辑过于复杂无法完全向量化，可通过Numba将循环函数编译为机器码，大幅提升速度：

from numba import jit

@jit(nopython=True)  # nopython模式性能最优
def numba_color_map(values, minimum=0, maximum=255):
    n = len(values)
    colors = np.zeros((n, 3), dtype=np.int32)
    minimum, maximum = float(minimum), float(maximum)
    for idx in range(n):
        val = values[idx]
        ratio = 2 * (val - minimum) / (maximum - minimum)
        b = max(0, int(255 * (1 - ratio)))
        r = max(0, int(255 * (ratio - 1)))
        g = 255 - b - r
        colors[idx] = (r, g, b)
    return colors

# 调用
points_colors = numba_color_map(points[:,3]) / 255.0

同样可给rainbow_color_map添加@jit(nopython=True)装饰器，注意提前将colors参数转为NumPy数组传入，适配Numba的索引规则。

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