Python 生成帶有RGBA信息的 ROS PointCloud2 message

起因

最近有一個工作是需要把一組三維點以ROS PointCloud2 messge的形式進行publish。並且需要使用python環境。原始點雲只有座標數據，需要根據點距離座標原點的距離對點雲進行上色。

經過

通過參考一些開源項目的源碼，並通過NumPy的structured array功能實現了PointCloud2 message的生成。其中需要注意的是一點，PointCloud2 message 通過一個uint32類型的整數表達RGBA四個通道。雖然稱作RGBA，但是在uint32的4個字節中，從高位向低位分別是B, G, R, A通道。另外，對2維的structured array進行索引，其方法和行爲可能都和普通的高維array有些差異。

結果

直接上源碼。

import numpy as np

import rospy
from sensor_msgs.msg import PointCloud2, PointField
from std_msgs.msg import Header
from ColorMapping import color_map

DIST_COLORS = [\
    "#2980b9",\
    "#27ae60",\
    "#f39c12",\
    "#c0392b",\
    ]

DIST_COLOR_LEVELS = 20

def convert_numpy_2_pointcloud2_color(points, stamp=None, frame_id=None, maxDistColor=None):
    '''
    Create a sensor_msgs.PointCloud2 from an array
    of points. 

    This function will automatically assign RGB values to each point. The RGB values are
    determined by the distance of a point from the origin. Use maxDistColor to set the distance 
    at which the color corresponds to the farthest distance is used.

    points: A NumPy array of Nx3.
    stamp: An alternative ROS header stamp.
    frame_id: The frame id. String.
    maxDisColor: Should be positive if specified..

    This function get inspired by 
    https://github.com/spillai/pybot/blob/master/pybot/externals/ros/pointclouds.py
    https://gist.github.com/lucasw/ea04dcd65bc944daea07612314d114bb
    (https://answers.ros.org/question/289576/understanding-the-bytes-in-a-pcl2-message/)
    and expo_utility.xyz_array_to_point_cloud_msg() function of the AirSim package.

    ROS sensor_msgs/PointField.
    http://docs.ros.org/melodic/api/sensor_msgs/html/msg/PointField.html

    More references on mixed-type NumPy array, structured array.
    https://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html
    https://stackoverflow.com/questions/37791134/merge-width-x-height-x-3-numpy-uint8-array-into-width-x-height-x-1-uint32-array
    https://jakevdp.github.io/PythonDataScienceHandbook/02.09-structured-data-numpy.html
    '''
    
    # Clipping input.
    dist = np.linalg.norm( points, axis=1 )
    if ( maxDistColor is not None and maxDistColor > 0):
        dist = np.clip(dist, 0, maxDistColor)

    # Compose color.
    cr, cg, cb = color_map( dist, DIST_COLORS, DIST_COLOR_LEVELS )

    C = np.zeros((cr.size, 4), dtype=np.uint8) + 255

    C[:, 0] = cb.astype(np.uint8)
    C[:, 1] = cg.astype(np.uint8)
    C[:, 2] = cr.astype(np.uint8)

    C = C.view("uint32")

    # Structured array.
    pointsColor = np.zeros( (points.shape[0], 1), \
        dtype={ 
            "names": ( "x", "y", "z", "rgba" ), 
            "formats": ( "f4", "f4", "f4", "u4" )} )

    points = points.astype(np.float32)

    pointsColor["x"] = points[:, 0].reshape((-1, 1))
    pointsColor["y"] = points[:, 1].reshape((-1, 1))
    pointsColor["z"] = points[:, 2].reshape((-1, 1))
    pointsColor["rgba"] = C

    header = Header()

    if stamp is None:
        header.stamp = rospy.Time().now()
    else:
        header.stamp = stamp

    if frame_id is None:
        header.frame_id = "None"
    else:
        header.frame_id = frame_id

    msg = PointCloud2()
    msg.header = header

    if len(points.shape) == 3:
        msg.height = points.shape[1]
        msg.width = points.shape[0]
    else:
        msg.height = 1
        msg.width  = points.shape[0]

    msg.fields = [
        PointField('x',  0, PointField.FLOAT32, 1),
        PointField('y',  4, PointField.FLOAT32, 1),
        PointField('z',  8, PointField.FLOAT32, 1),
        PointField('rgb', 12, PointField.UINT32, 1),
        ]

    msg.is_bigendian = False
    msg.point_step   = 16
    msg.row_step     = msg.point_step * points.shape[0]
    msg.is_dense     = int( np.isfinite(points).all() )
    msg.data         = pointsColor.tostring()

    return msg

其中，ColorMapping是我參考另外一個博客，在其源碼基礎上自己創建的一個生成線性梯度RGB值的輔助工具。源碼在這裏。

參考

ROS PointCloud2 message生成
https://github.com/spillai/pybot/blob/master/pybot/externals/ros/pointclouds.py https://gist.github.com/lucasw/ea04dcd65bc944daea07612314d114bb
(https://answers.ros.org/question/289576/understanding-the-bytes-in-a-pcl2-message/)

ROS sensor_msgs/PointField
http://docs.ros.org/melodic/api/sensor_msgs/html/msg/PointField.html

NumPy structured array
https://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html
https://stackoverflow.com/questions/37791134/merge-width-x-height-x-3-numpy-uint8-array-into-width-x-height-x-1-uint32-array
https://jakevdp.github.io/PythonDataScienceHandbook/02.09-structured-data-numpy.html

Color mapping
https://bsou.io/posts/color-gradients-with-python