point_cloud2.py
Go to the documentation of this file.
00001 # Software License Agreement (BSD License)
00002 #
00003 # Copyright (c) 2008, Willow Garage, Inc.
00004 # All rights reserved.
00005 #
00006 # Redistribution and use in source and binary forms, with or without
00007 # modification, are permitted provided that the following conditions
00008 # are met:
00009 #
00010 #  * Redistributions of source code must retain the above copyright
00011 #    notice, this list of conditions and the following disclaimer.
00012 #  * Redistributions in binary form must reproduce the above
00013 #    copyright notice, this list of conditions and the following
00014 #    disclaimer in the documentation and/or other materials provided
00015 #    with the distribution.
00016 #  * Neither the name of Willow Garage, Inc. nor the names of its
00017 #    contributors may be used to endorse or promote products derived
00018 #    from this software without specific prior written permission.
00019 #
00020 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
00021 # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
00022 # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
00023 # FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
00024 # COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
00025 # INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
00026 # BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
00027 # LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
00028 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
00029 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
00030 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
00031 # POSSIBILITY OF SUCH DAMAGE.
00032 #
00033 # Author: Jon Binney
00034 '''
00035 Functions for working with PointCloud2.
00036 '''
00037 
00038 __docformat__ = "restructuredtext en"
00039 
00040 from .registry import converts_from_numpy, converts_to_numpy
00041 
00042 import numpy as np
00043 from sensor_msgs.msg import PointCloud2, PointField
00044 
00045 # prefix to the names of dummy fields we add to get byte alignment correct. this needs to not
00046 # clash with any actual field names
00047 DUMMY_FIELD_PREFIX = '__'
00048 
00049 # mappings between PointField types and numpy types
00050 type_mappings = [(PointField.INT8, np.dtype('int8')), (PointField.UINT8, np.dtype('uint8')), (PointField.INT16, np.dtype('int16')),
00051                  (PointField.UINT16, np.dtype('uint16')), (PointField.INT32, np.dtype('int32')), (PointField.UINT32, np.dtype('uint32')),
00052                  (PointField.FLOAT32, np.dtype('float32')), (PointField.FLOAT64, np.dtype('float64'))]
00053 pftype_to_nptype = dict(type_mappings)
00054 nptype_to_pftype = dict((nptype, pftype) for pftype, nptype in type_mappings)
00055 
00056 # sizes (in bytes) of PointField types
00057 pftype_sizes = {PointField.INT8: 1, PointField.UINT8: 1, PointField.INT16: 2, PointField.UINT16: 2,
00058                 PointField.INT32: 4, PointField.UINT32: 4, PointField.FLOAT32: 4, PointField.FLOAT64: 8}
00059 
00060 @converts_to_numpy(PointField, plural=True)
00061 def fields_to_dtype(fields, point_step):
00062     '''Convert a list of PointFields to a numpy record datatype.
00063     '''
00064     offset = 0
00065     np_dtype_list = []
00066     for f in fields:
00067         while offset < f.offset:
00068             # might be extra padding between fields
00069             np_dtype_list.append(('%s%d' % (DUMMY_FIELD_PREFIX, offset), np.uint8))
00070             offset += 1
00071 
00072         dtype = pftype_to_nptype[f.datatype]
00073         if f.count != 1:
00074             dtype = np.dtype((dtype, f.count))
00075 
00076         np_dtype_list.append((f.name, dtype))
00077         offset += pftype_sizes[f.datatype] * f.count
00078 
00079     # might be extra padding between points
00080     while offset < point_step:
00081         np_dtype_list.append(('%s%d' % (DUMMY_FIELD_PREFIX, offset), np.uint8))
00082         offset += 1
00083         
00084     return np_dtype_list
00085 
00086 
00087 @converts_from_numpy(PointField, plural=True)
00088 def dtype_to_fields(dtype):
00089     '''Convert a numpy record datatype into a list of PointFields.
00090     '''
00091     fields = []
00092     for field_name in dtype.names:
00093         np_field_type, field_offset = dtype.fields[field_name]
00094         pf = PointField()
00095         pf.name = field_name
00096         if np_field_type.subdtype:
00097             item_dtype, shape = np_field_type.subdtype
00098             pf.count = np.prod(shape)
00099             np_field_type = item_dtype
00100         else:
00101             pf.count = 1
00102 
00103         pf.datatype = nptype_to_pftype[np_field_type]
00104         pf.offset = field_offset
00105         fields.append(pf)
00106     return fields
00107 
00108 @converts_to_numpy(PointCloud2)
00109 def pointcloud2_to_array(cloud_msg, squeeze=True):
00110     ''' Converts a rospy PointCloud2 message to a numpy recordarray 
00111     
00112     Reshapes the returned array to have shape (height, width), even if the height is 1.
00113 
00114     The reason for using np.frombuffer rather than struct.unpack is speed... especially
00115     for large point clouds, this will be <much> faster.
00116     '''
00117     # construct a numpy record type equivalent to the point type of this cloud
00118     dtype_list = fields_to_dtype(cloud_msg.fields, cloud_msg.point_step)
00119 
00120     # parse the cloud into an array
00121     cloud_arr = np.frombuffer(cloud_msg.data, dtype_list)
00122 
00123     # remove the dummy fields that were added
00124     cloud_arr = cloud_arr[
00125         [fname for fname, _type in dtype_list if not (fname[:len(DUMMY_FIELD_PREFIX)] == DUMMY_FIELD_PREFIX)]]
00126     
00127     if squeeze and cloud_msg.height == 1:
00128         return np.reshape(cloud_arr, (cloud_msg.width,))
00129     else:
00130         return np.reshape(cloud_arr, (cloud_msg.height, cloud_msg.width))
00131 
00132 @converts_from_numpy(PointCloud2)
00133 def array_to_pointcloud2(cloud_arr, stamp=None, frame_id=None):
00134     '''Converts a numpy record array to a sensor_msgs.msg.PointCloud2.
00135     '''
00136     # make it 2d (even if height will be 1)
00137     cloud_arr = np.atleast_2d(cloud_arr)
00138 
00139     cloud_msg = PointCloud2()
00140 
00141     if stamp is not None:
00142         cloud_msg.header.stamp = stamp
00143     if frame_id is not None:
00144         cloud_msg.header.frame_id = frame_id
00145     cloud_msg.height = cloud_arr.shape[0]
00146     cloud_msg.width = cloud_arr.shape[1]
00147     cloud_msg.fields = dtype_to_fields(cloud_arr.dtype)
00148     cloud_msg.is_bigendian = False # assumption
00149     cloud_msg.point_step = cloud_arr.dtype.itemsize
00150     cloud_msg.row_step = cloud_msg.point_step*cloud_arr.shape[1]
00151     cloud_msg.is_dense = all([np.isfinite(cloud_arr[fname]).all() for fname in cloud_arr.dtype.names])
00152     cloud_msg.data = cloud_arr.tostring()
00153     return cloud_msg
00154 
00155 def merge_rgb_fields(cloud_arr):
00156     '''Takes an array with named np.uint8 fields 'r', 'g', and 'b', and returns an array in
00157     which they have been merged into a single np.float32 'rgb' field. The first byte of this
00158     field is the 'r' uint8, the second is the 'g', uint8, and the third is the 'b' uint8.
00159 
00160     This is the way that pcl likes to handle RGB colors for some reason.
00161     '''
00162     r = np.asarray(cloud_arr['r'], dtype=np.uint32)
00163     g = np.asarray(cloud_arr['g'], dtype=np.uint32)
00164     b = np.asarray(cloud_arr['b'], dtype=np.uint32)    
00165     rgb_arr = np.array((r << 16) | (g << 8) | (b << 0), dtype=np.uint32)
00166 
00167     # not sure if there is a better way to do this. i'm changing the type of the array
00168     # from uint32 to float32, but i don't want any conversion to take place -jdb
00169     rgb_arr.dtype = np.float32
00170 
00171     # create a new array, without r, g, and b, but with rgb float32 field
00172     new_dtype = []
00173     for field_name in cloud_arr.dtype.names:
00174         field_type, field_offset = cloud_arr.dtype.fields[field_name]
00175         if field_name not in ('r', 'g', 'b'):
00176             new_dtype.append((field_name, field_type))
00177     new_dtype.append(('rgb', np.float32))
00178     new_cloud_arr = np.zeros(cloud_arr.shape, new_dtype)    
00179 
00180     # fill in the new array
00181     for field_name in new_cloud_arr.dtype.names:
00182         if field_name == 'rgb':
00183             new_cloud_arr[field_name] = rgb_arr
00184         else:
00185             new_cloud_arr[field_name] = cloud_arr[field_name]
00186         
00187     return new_cloud_arr
00188 
00189 def split_rgb_field(cloud_arr):
00190     '''Takes an array with a named 'rgb' float32 field, and returns an array in which
00191     this has been split into 3 uint 8 fields: 'r', 'g', and 'b'.
00192 
00193     (pcl stores rgb in packed 32 bit floats)
00194     '''
00195     rgb_arr = cloud_arr['rgb'].copy()
00196     rgb_arr.dtype = np.uint32
00197     r = np.asarray((rgb_arr >> 16) & 255, dtype=np.uint8)
00198     g = np.asarray((rgb_arr >> 8) & 255, dtype=np.uint8)
00199     b = np.asarray(rgb_arr & 255, dtype=np.uint8)
00200     
00201     # create a new array, without rgb, but with r, g, and b fields
00202     new_dtype = []
00203     for field_name in cloud_arr.dtype.names:
00204         field_type, field_offset = cloud_arr.dtype.fields[field_name]
00205         if not field_name == 'rgb':
00206             new_dtype.append((field_name, field_type))
00207     new_dtype.append(('r', np.uint8))
00208     new_dtype.append(('g', np.uint8))
00209     new_dtype.append(('b', np.uint8))    
00210     new_cloud_arr = np.zeros(cloud_arr.shape, new_dtype)
00211     
00212     # fill in the new array
00213     for field_name in new_cloud_arr.dtype.names:
00214         if field_name == 'r':
00215             new_cloud_arr[field_name] = r
00216         elif field_name == 'g':
00217             new_cloud_arr[field_name] = g
00218         elif field_name == 'b':
00219             new_cloud_arr[field_name] = b
00220         else:
00221             new_cloud_arr[field_name] = cloud_arr[field_name]
00222     return new_cloud_arr
00223 
00224 def get_xyz_points(cloud_array, remove_nans=True, dtype=np.float):
00225     '''Pulls out x, y, and z columns from the cloud recordarray, and returns
00226         a 3xN matrix.
00227     '''
00228     # remove crap points
00229     if remove_nans:
00230         mask = np.isfinite(cloud_array['x']) & np.isfinite(cloud_array['y']) & np.isfinite(cloud_array['z'])
00231         cloud_array = cloud_array[mask]
00232     
00233     # pull out x, y, and z values
00234     points = np.zeros(cloud_array.shape + (3,), dtype=dtype)
00235     points[...,0] = cloud_array['x']
00236     points[...,1] = cloud_array['y']
00237     points[...,2] = cloud_array['z']
00238 
00239     return points
00240 
00241 def pointcloud2_to_xyz_array(cloud_msg, remove_nans=True):
00242     return get_xyz_points(pointcloud2_to_array(cloud_msg), remove_nans=remove_nans)


ros_numpy
Author(s): Eric Wieser
autogenerated on Wed Apr 3 2019 02:40:14