Add support for CompressedImage, XXXStamped and messages from rosbag

README.md

@@ -26,12 +26,15 @@ Currently supports:
    ```
 
 * `sensor_msgs.msg.Image` ↔ 2/3-D `np.array`, similar to the function of `cv_bridge`, but without the dependency on `cv2`
+* `sensor_msgs.msg.CompressedImage` → `np.array`. 
 * `nav_msgs.msg.OccupancyGrid` ↔ `np.ma.array`
-* `geometry.msg.Vector3` ↔ 1-D `np.array`. `hom=True` gives `[x, y, z, 0]`
-* `geometry.msg.Point` ↔ 1-D `np.array`. `hom=True` gives `[x, y, z, 1]`
-* `geometry.msg.Quaternion` ↔ 1-D `np.array`, `[x, y, z, w]`
-* `geometry.msg.Transform` ↔ 4×4 `np.array`, the homogeneous transformation matrix
-* `geometry.msg.Pose` ↔ 4×4 `np.array`, the homogeneous transformation matrix from the origin
+* `geometry_msg.Vector3` ↔ 1-D `np.array`. `hom=True` gives `[x, y, z, 0]`
+* `geometry_msg.Point` ↔ 1-D `np.array`. `hom=True` gives `[x, y, z, 1]`
+* `geometry_msg.Quaternion` ↔ 1-D `np.array`, `[x, y, z, w]`
+* `geometry_msg.Transform` ↔ 4×4 `np.array`, the homogeneous transformation matrix
+* `geometry_msg.Pose` ↔ 4×4 `np.array`, the homogeneous transformation matrix from the origin
+* `geometry_msg.XXXXXStamped` → `np.array`.
+* All the previous types but directly taken from a rosbag (meanwhile you have the same version of the messages, i.e. same md5sum).
 
 Support for more types can be added with:
 

src/ros_numpy/geometry.py

@@ -1,5 +1,6 @@
 from .registry import converts_from_numpy, converts_to_numpy
 from geometry_msgs.msg import Transform, Vector3, Quaternion, Point, Pose
+from geometry_msgs.msg import TransformStamped, Vector3Stamped, QuaternionStamped, PointStamped, PoseStamped
 from . import numpify
 
 import numpy as np
@@ -13,6 +14,10 @@ def vector3_to_numpy(msg, hom=False):
 	else:
 		return np.array([msg.x, msg.y, msg.z])
 
+@converts_to_numpy(Vector3Stamped)
+def vector3stamped_to_numpy(msg, hom=False):
+	return vector3_to_numpy(msg.vector, hom)
+
 @converts_from_numpy(Vector3)
 def numpy_to_vector3(arr):
 	if arr.shape[-1] == 4:
@@ -31,6 +36,10 @@ def point_to_numpy(msg, hom=False):
 	else:
 		return np.array([msg.x, msg.y, msg.z])
 
+@converts_to_numpy(PointStamped)
+def pointstamped_to_numpy(msg, hom=False):
+	return point_to_numpy(msg.point, hom)
+
 @converts_from_numpy(Point)
 def numpy_to_point(arr):
 	if arr.shape[-1] == 4:
@@ -45,6 +54,10 @@ def numpy_to_point(arr):
 def quat_to_numpy(msg):
 	return np.array([msg.x, msg.y, msg.z, msg.w])
 
+@converts_to_numpy(QuaternionStamped)
+def quatstamped_to_numpy(msg):
+	return quat_to_numpy(msg.quaternion)
+
 @converts_from_numpy(Quaternion)
 def numpy_to_quat(arr):
 	assert arr.shape[-1] == 4
@@ -67,6 +80,10 @@ def transform_to_numpy(msg):
         transformations.quaternion_matrix(numpify(msg.rotation))
     )
 
+@converts_to_numpy(TransformStamped)
+def transformstamped_to_numpy(msg):
+	return transform_to_numpy(msg.transform)
+
 @converts_from_numpy(Transform)
 def numpy_to_transform(arr):
 	from tf import transformations
@@ -99,6 +116,10 @@ def pose_to_numpy(msg):
         transformations.quaternion_matrix(numpify(msg.orientation))
     )
 
+@converts_to_numpy(PoseStamped)
+def posestamped_to_numpy(msg):
+	return pose_to_numpy(msg.pose)
+
 @converts_from_numpy(Pose)
 def numpy_to_pose(arr):
 	from tf import transformations

src/ros_numpy/image.py

@@ -1,7 +1,7 @@
 import sys
 
 from .registry import converts_from_numpy, converts_to_numpy
-from sensor_msgs.msg import Image
+from sensor_msgs.msg import Image, CompressedImage
 
 import numpy as np
 from numpy.lib.stride_tricks import as_strided
@@ -80,6 +80,9 @@ def image_to_numpy(msg):
 		data = data[...,0]
 	return data
 
+@converts_to_numpy(CompressedImage)
+def compressed_image_to_numpy(msg):
+	return np.fromstring(msg.data, np.uint8)
 
 @converts_from_numpy(Image)
 def numpy_to_image(arr, encoding):

src/ros_numpy/registry.py

@@ -32,6 +32,21 @@ def numpify(msg, *args, **kwargs):
 			raise ValueError("Cannot determine the type of an empty Collection")
 		conv = _to_numpy.get((msg[0].__class__, True))
 
+	if not conv:
+		# on messages coming from a rosbag, __class__ instead of
+		# looking like '_sensor_msgs__Image' looks like tmpldS_vh._sensor_msgs__Image
+		# so we try to check for the _type instead ('sensor_msgs/Image')
+		for class_instance, _ in _to_numpy.keys():
+			if msg._type == class_instance._type:
+				# we make sure it's the same as the installed message
+				if msg._md5sum == class_instance._md5sum:
+					conv = _to_numpy.get((class_instance, False))
+					break
+				else:
+					raise ValueError("Unable to convert message of type {}, ".format(msg._type) +
+					 "the md5sum of message to numpify and the installed message differ: {} != {}".format(
+						msg._md5sum, class_instance._md5sum)
+					)
 
 	if not conv:
 		raise ValueError("Unable to convert message {} - only supports {}".format(

test/create_test_bag.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python
+# Create a rosbag with the supported types
+import rosbag
+from ros_numpy import msgify, numpy_msg
+import numpy as np
+from sensor_msgs.msg import Image, PointCloud2, CompressedImage
+from nav_msgs.msg import OccupancyGrid
+from geometry_msgs.msg import Vector3, Point, Quaternion, Transform, Pose
+import zlib, struct
+
+
+def makeArray(npoints):
+	points_arr = np.zeros((npoints,), dtype=[
+		('x', np.float32),
+		('y', np.float32),
+		('z', np.float32),
+		('r', np.uint8),
+		('g', np.uint8),
+		('b', np.uint8)])
+	points_arr['x'] = np.random.random((npoints,))
+	points_arr['y'] = np.random.random((npoints,))
+	points_arr['z'] = np.random.random((npoints,))
+	points_arr['r'] = np.floor(np.random.random((npoints,)) * 255)
+	points_arr['g'] = 0
+	points_arr['b'] = 255
+
+	return points_arr
+
+# From http://stackoverflow.com/questions/902761/saving-a-numpy-array-as-an-image
+def write_png(buf, width, height):
+	""" buf: must be bytes or a bytearray in Python3.x,
+		a regular string in Python2.x.
+	"""
+
+	# reverse the vertical line order and add null bytes at the start
+	width_byte_4 = width * 4
+	raw_data = b''.join(b'\x00' + buf[span:span + width_byte_4]
+						for span in range((height - 1) * width_byte_4, -1, - width_byte_4))
+
+	def png_pack(png_tag, data):
+		chunk_head = png_tag + data
+		return (struct.pack("!I", len(data)) +
+				chunk_head +
+				struct.pack("!I", 0xFFFFFFFF & zlib.crc32(chunk_head)))
+
+	return b''.join([
+		b'\x89PNG\r\n\x1a\n',
+		png_pack(b'IHDR', struct.pack("!2I5B", width, height, 8, 6, 0, 0, 0)),
+		png_pack(b'IDAT', zlib.compress(raw_data, 9)),
+		png_pack(b'IEND', b'')])
+
+def get_png_numpy_array(array):
+	if any([len(row) != len(array[0]) for row in array]):
+		raise ValueError, "Array should have elements of equal size"
+
+		# First row becomes top row of image.
+	flat = []
+	map(flat.extend, reversed(array))
+	# Big-endian, unsigned 32-byte integer.
+	buf = b''.join([struct.pack('>I', ((0xffFFff & i32) << 8) | (i32 >> 24))
+					for i32 in flat])  # Rotate from ARGB to RGBA.
+
+	data = write_png(buf, len(array[0]), len(array))
+	return data
+
+bag = rosbag.Bag('test.bag', 'w')
+
+try:
+	arr = np.random.randint(0, 256, size=(240, 360, 3)).astype(np.uint8)
+	img = msgify(Image, arr, encoding='rgb8')
+	bag.write('/image', img)
+
+	points_arr = makeArray(100)
+	cloud_msg = msgify(numpy_msg(PointCloud2), points_arr)
+	bag.write('/pointcloud', cloud_msg)
+
+	data = -np.ones((30, 30), np.int8)
+	data[10:20, 10:20] = 100
+	occgrid = msgify(OccupancyGrid, data)
+	bag.write('/occupancygrid', occgrid)
+
+	v = Vector3(1.0, 2.0, 3.0)
+	bag.write('/vector3', v)
+
+	p = Point(0.0, 0.1, 0.3)
+	bag.write('/point', p)
+
+	q = Quaternion(0.0, 0.0, 0.0, 1.0)
+	bag.write('/quaternion', q)
+
+	t = Transform(v, q)
+	bag.write('/transform', t)
+
+	ps = Pose(p, q)
+	bag.write('/pose', ps)
+
+	ci = CompressedImage()
+	ci.format = 'png'
+	ci.data = get_png_numpy_array([[0xffFF0000, 0xffFFFF00],
+								   [0xff00aa77, 0xff333333]])
+	bag.write('/compressedimage', ci)
+
+finally:
+	bag.close()

test/different_md5_Pose.py

@@ -0,0 +1,135 @@
+# This Python file uses the following encoding: utf-8
+"""Taken from
+ros-indigo-geometry-msgs 1.11.9-0trusty-20160627-170848-0700
+and modified to have a made up _md5sum for
+testing purposes"""
+import sys
+python3 = True if sys.hexversion > 0x03000000 else False
+import genpy
+import struct
+
+import geometry_msgs.msg
+
+class Pose(genpy.Message):
+  _md5sum = "13371337133713371337133713371337"
+  _type = "geometry_msgs/Pose"
+  _has_header = False #flag to mark the presence of a Header object
+  _full_text = """# A representation of pose in free space, composed of postion and orientation. 
+Point position
+Quaternion orientation
+
+================================================================================
+MSG: geometry_msgs/Point
+# This contains the position of a point in free space
+float64 x
+float64 y
+float64 z
+
+================================================================================
+MSG: geometry_msgs/Quaternion
+# This represents an orientation in free space in quaternion form.
+
+float64 x
+float64 y
+float64 z
+float64 w
+"""
+  __slots__ = ['position','orientation']
+  _slot_types = ['geometry_msgs/Point','geometry_msgs/Quaternion']
+
+  def __init__(self, *args, **kwds):
+    """
+    Constructor. Any message fields that are implicitly/explicitly
+    set to None will be assigned a default value. The recommend
+    use is keyword arguments as this is more robust to future message
+    changes.  You cannot mix in-order arguments and keyword arguments.
+
+    The available fields are:
+       position,orientation
+
+    :param args: complete set of field values, in .msg order
+    :param kwds: use keyword arguments corresponding to message field names
+    to set specific fields.
+    """
+    if args or kwds:
+      super(Pose, self).__init__(*args, **kwds)
+      #message fields cannot be None, assign default values for those that are
+      if self.position is None:
+        self.position = geometry_msgs.msg.Point()
+      if self.orientation is None:
+        self.orientation = geometry_msgs.msg.Quaternion()
+    else:
+      self.position = geometry_msgs.msg.Point()
+      self.orientation = geometry_msgs.msg.Quaternion()
+
+  def _get_types(self):
+    """
+    internal API method
+    """
+    return self._slot_types
+
+  def serialize(self, buff):
+    """
+    serialize message into buffer
+    :param buff: buffer, ``StringIO``
+    """
+    try:
+      _x = self
+      buff.write(_struct_7d.pack(_x.position.x, _x.position.y, _x.position.z, _x.orientation.x, _x.orientation.y, _x.orientation.z, _x.orientation.w))
+    except struct.error as se: self._check_types(struct.error("%s: '%s' when writing '%s'" % (type(se), str(se), str(locals().get('_x', self)))))
+    except TypeError as te: self._check_types(ValueError("%s: '%s' when writing '%s'" % (type(te), str(te), str(locals().get('_x', self)))))
+
+  def deserialize(self, str):
+    """
+    unpack serialized message in str into this message instance
+    :param str: byte array of serialized message, ``str``
+    """
+    try:
+      if self.position is None:
+        self.position = geometry_msgs.msg.Point()
+      if self.orientation is None:
+        self.orientation = geometry_msgs.msg.Quaternion()
+      end = 0
+      _x = self
+      start = end
+      end += 56
+      (_x.position.x, _x.position.y, _x.position.z, _x.orientation.x, _x.orientation.y, _x.orientation.z, _x.orientation.w,) = _struct_7d.unpack(str[start:end])
+      return self
+    except struct.error as e:
+      raise genpy.DeserializationError(e) #most likely buffer underfill
+
+
+  def serialize_numpy(self, buff, numpy):
+    """
+    serialize message with numpy array types into buffer
+    :param buff: buffer, ``StringIO``
+    :param numpy: numpy python module
+    """
+    try:
+      _x = self
+      buff.write(_struct_7d.pack(_x.position.x, _x.position.y, _x.position.z, _x.orientation.x, _x.orientation.y, _x.orientation.z, _x.orientation.w))
+    except struct.error as se: self._check_types(struct.error("%s: '%s' when writing '%s'" % (type(se), str(se), str(locals().get('_x', self)))))
+    except TypeError as te: self._check_types(ValueError("%s: '%s' when writing '%s'" % (type(te), str(te), str(locals().get('_x', self)))))
+
+  def deserialize_numpy(self, str, numpy):
+    """
+    unpack serialized message in str into this message instance using numpy for array types
+    :param str: byte array of serialized message, ``str``
+    :param numpy: numpy python module
+    """
+    try:
+      if self.position is None:
+        self.position = geometry_msgs.msg.Point()
+      if self.orientation is None:
+        self.orientation = geometry_msgs.msg.Quaternion()
+      end = 0
+      _x = self
+      start = end
+      end += 56
+      (_x.position.x, _x.position.y, _x.position.z, _x.orientation.x, _x.orientation.y, _x.orientation.z, _x.orientation.w,) = _struct_7d.unpack(str[start:end])
+      return self
+    except struct.error as e:
+      raise genpy.DeserializationError(e) #most likely buffer underfill
+
+_struct_I = genpy.struct_I
+_struct_7d = struct.Struct("<7d")