ROS Rviz custom Point Cloud visualization example

point_cloud_visualization.py

#!/usr/bin/env python

"""
Rviz point cloud visualization marker example
Author: Roberto Zegers R.
Copyright: Copyright (c) 2021, Roberto Zegers R.
License: BSD-3-Clause
Date: March 2021
"""

import rospy
from visualization_msgs.msg import MarkerArray, Marker
from std_msgs.msg import  Header, ColorRGBA
from geometry_msgs.msg import Point, Vector3
from sensor_msgs import point_cloud2
from sensor_msgs.msg import PointCloud2, PointField
from std_msgs.msg import Header

import struct
from random import randrange as rand

class GridVisualizer:
    def __init__(self, flat_map, resolution, origin, width, id=0, frame="map"):
        self.flat_map = flat_map
        self.map_resolution = resolution
        self.map_origin = origin
        self.map_width = width
        self.id = id
        self.rgba_colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']
        self.plot_cloud = rospy.Publisher('/grid_viz', PointCloud2, queue_size=1)
        self.header = Header()
        self.header.frame_id = frame
        # field name, offset, datatype(float32), count
        self.field_x = PointField('x', 0, PointField.FLOAT32, 1)
        self.field_y = PointField('y', 4, PointField.FLOAT32, 1)
        self.field_z = PointField('z', 8, PointField.FLOAT32, 1)
        self.field_rgba = PointField('rgba', 12, PointField.UINT32, 1)
        self.fields = [self.field_x, self.field_y, self.field_z, self.field_rgba]
        self.points = []
        self.init_points()

    def init_points(self):
        # set RGBA color in 0-255 range, including transparency
        r, g, b, a = 0, 255, 255, 255 # ints!
        rgba_to_hex = struct.unpack('I', struct.pack('BBBB', b, g, r, a))[0]
        for idx, value in enumerate(self.flat_map):
            xyz_point = self.indexToWorld(idx)
            # cloud points initialized to RGBA color defined above
            self.points.append(xyz_point + [rgba_to_hex]) # [x, y, z, rgba]

    def set_color(self, idx, rgba_list):
        r, g, b, a = rgba_list
        rgba_to_hex = struct.unpack('I', struct.pack('BBBB', b, g, r, a))[0]
        self.points[idx][3] = rgba_to_hex

        pc2 = point_cloud2.create_cloud(self.header, self.fields, self.points)
        pc2.header.stamp = rospy.Time.now()
        self.plot_cloud.publish(pc2)

    def indexToWorld(self, array_index):
        """
        Converts a linear index value to world coordinates (meters) in the form of a list
        array_index: a linear index value, specifying a cell/pixel in an 1-D array
        returns: [x,y,z] list
        """
        # convert to x,y grid cell/pixel coordinates
        grid_cell_x = array_index % self.map_width # modulo yields the remainder
        grid_cell_y = array_index // self.map_width # floor division '//' yields the quotient
        # convert to world coordinates
        x = self.map_resolution * grid_cell_x + self.map_origin[0] + self.map_resolution / 2
        y = self.map_resolution * grid_cell_y + self.map_origin[1] + self.map_resolution / 2
        z = 0.0
        return [x,y,z]

def main():
    try:
        map_size = 10000
        # map as a 1-d array 
        fake_map = [x for x in range(map_size)]
        map_resolution = 0.2
        # where the bottom left pixel of the point cloud is located relative to the maps origin in meters
        # To center the cloud at the maps center take width/2 * resolution
        origin = [-10,-10]
        width = 100
        custom_point_cloud = GridVisualizer(fake_map, map_resolution, origin, width)

        # Infinitely select a random point and change its color (at a rate)
        rgba_list = [255, 0, 0, 255]
        while not rospy.is_shutdown():
          # Select a random point anywhere in the point cloud
          random_point = rand(map_size)
          custom_point_cloud.set_color(random_point, rgba_list)
          rate.sleep()

    except rospy.ROSInterruptException:
        pass

if __name__ == '__main__':
    rospy.init_node("cloud_test")
    rate = rospy.Rate(1)
    main()