Skip to content

Instantly share code, notes, and snippets.

@adioshun
Created June 4, 2018 07:02
Show Gist options
  • Save adioshun/f35919c895631314394aa1762c24334c to your computer and use it in GitHub Desktop.
Save adioshun/f35919c895631314394aa1762c24334c to your computer and use it in GitHub Desktop.
#!/usr/bin/env python
# Copyright (C) 2017 Electric Movement Inc.
#
# This file is part of Robotic Arm: Pick and Place project for Udacity
# Robotics nano-degree program
#
# All Rights Reserved.
# Author: Harsh Pandya
#
# Import modules
import rospy
import pcl
import numpy as np
import ctypes
import struct
import sensor_msgs.point_cloud2 as pc2
from sensor_msgs.msg import PointCloud2, PointField
from std_msgs.msg import Header
from random import randint
def random_color_gen():
""" Generates a random color
Args: None
Returns:
list: 3 elements, R, G, and B
"""
r = randint(0, 255)
g = randint(0, 255)
b = randint(0, 255)
return [r, g, b]
def ros_to_pcl(ros_cloud):
""" Converts a ROS PointCloud2 message to a pcl PointXYZRGB
Args:
ros_cloud (PointCloud2): ROS PointCloud2 message
Returns:
pcl.PointCloud_PointXYZRGB: PCL XYZRGB point cloud
"""
points_list = []
for data in pc2.read_points(ros_cloud, skip_nans=True):
points_list.append([data[0], data[1], data[2], data[3]])
pcl_data = pcl.PointCloud_PointXYZRGB()
pcl_data.from_list(points_list)
return pcl_data
def pcl_to_ros(pcl_array):
""" Converts a pcl PointXYZRGB to a ROS PointCloud2 message
Args:
pcl_array (PointCloud_PointXYZRGB): A PCL XYZRGB point cloud
Returns:
PointCloud2: A ROS point cloud
"""
ros_msg = PointCloud2()
ros_msg.header.stamp = rospy.Time.now()
ros_msg.header.frame_id = "world"
ros_msg.height = 1
ros_msg.width = pcl_array.size
ros_msg.fields.append(PointField(
name="x",
offset=0,
datatype=PointField.FLOAT32, count=1))
ros_msg.fields.append(PointField(
name="y",
offset=4,
datatype=PointField.FLOAT32, count=1))
ros_msg.fields.append(PointField(
name="z",
offset=8,
datatype=PointField.FLOAT32, count=1))
ros_msg.fields.append(PointField(
name="rgb",
offset=16,
datatype=PointField.FLOAT32, count=1))
ros_msg.is_bigendian = False
ros_msg.point_step = 32
ros_msg.row_step = ros_msg.point_step * ros_msg.width * ros_msg.height
ros_msg.is_dense = False
buffer = []
for data in pcl_array:
s = struct.pack('>f', data[3])
i = struct.unpack('>l', s)[0]
pack = ctypes.c_uint32(i).value
r = (pack & 0x00FF0000) >> 16
g = (pack & 0x0000FF00) >> 8
b = (pack & 0x000000FF)
buffer.append(struct.pack('ffffBBBBIII', data[0], data[1], data[2], 1.0, b, g, r, 0, 0, 0, 0))
ros_msg.data = "".join(buffer)
return ros_msg
def XYZRGB_to_XYZ(XYZRGB_cloud):
""" Converts a PCL XYZRGB point cloud to an XYZ point cloud (removes color info)
Args:
XYZRGB_cloud (PointCloud_PointXYZRGB): A PCL XYZRGB point cloud
Returns:
PointCloud_PointXYZ: A PCL XYZ point cloud
"""
XYZ_cloud = pcl.PointCloud()
points_list = []
for data in XYZRGB_cloud:
points_list.append([data[0], data[1], data[2]])
XYZ_cloud.from_list(points_list)
return XYZ_cloud
def XYZ_to_XYZRGB(XYZ_cloud, color):
""" Converts a PCL XYZ point cloud to a PCL XYZRGB point cloud
All returned points in the XYZRGB cloud will be the color indicated
by the color parameter.
Args:
XYZ_cloud (PointCloud_XYZ): A PCL XYZ point cloud
color (list): 3-element list of integers [0-255,0-255,0-255]
Returns:
PointCloud_PointXYZRGB: A PCL XYZRGB point cloud
"""
XYZRGB_cloud = pcl.PointCloud_PointXYZRGB()
points_list = []
float_rgb = rgb_to_float(color)
for data in XYZ_cloud:
points_list.append([data[0], data[1], data[2], float_rgb])
XYZRGB_cloud.from_list(points_list)
return XYZRGB_cloud
def rgb_to_float(color):
""" Converts an RGB list to the packed float format used by PCL
From the PCL docs:
"Due to historical reasons (PCL was first developed as a ROS package),
the RGB information is packed into an integer and casted to a float"
Args:
color (list): 3-element list of integers [0-255,0-255,0-255]
Returns:
float_rgb: RGB value packed as a float
"""
hex_r = (0xff & color[0]) << 16
hex_g = (0xff & color[1]) << 8
hex_b = (0xff & color[2])
hex_rgb = hex_r | hex_g | hex_b
float_rgb = struct.unpack('f', struct.pack('i', hex_rgb))[0]
return float_rgb
def float_to_rgb(float_rgb):
""" Converts a packed float RGB format to an RGB list
Args:
float_rgb: RGB value packed as a float
Returns:
color (list): 3-element list of integers [0-255,0-255,0-255]
"""
s = struct.pack('>f', float_rgb)
i = struct.unpack('>l', s)[0]
pack = ctypes.c_uint32(i).value
r = (pack & 0x00FF0000) >> 16
g = (pack & 0x0000FF00) >> 8
b = (pack & 0x000000FF)
color = [r,g,b]
return color
def get_color_list(cluster_count):
""" Returns a list of randomized colors
Args:
cluster_count (int): Number of random colors to generate
Returns:
(list): List containing 3-element color lists
"""
if (cluster_count > len(get_color_list.color_list)):
for i in xrange(len(get_color_list.color_list), cluster_count):
get_color_list.color_list.append(random_color_gen())
return get_color_list.color_list
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment