python scripts for SVM obj identification

scripts/pcl_helper.py

@@ -0,0 +1,217 @@
+#!/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 ROS PointCloud2 message to a pcl PointXYZRGB
+    
+        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