Add support for CompressedImage, XXXStamped and messages from rosbag

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")

test/test_from_bag.py

@@ -0,0 +1,136 @@
+import unittest
+import numpy as np
+import ros_numpy
+from rosbag import Bag
+from ros_numpy import msgify
+from sensor_msgs.msg import Image, PointCloud2, CompressedImage
+from nav_msgs.msg import OccupancyGrid
+from geometry_msgs.msg import Vector3, Point, Quaternion, Transform, Pose
+
+
+class TestImagesFromRosbag(unittest.TestCase):
+	def test_image_from_bag(self):
+		bag = Bag('test.bag')
+		# There is only one message
+		for topic, bagmsg, time in bag.read_messages(topics=['/image']):
+			arr = ros_numpy.numpify(bagmsg)
+			msg = Image()
+			for slot in msg.__slots__:
+				# Go thru all the slots, data, header...
+				setattr(msg, slot, getattr(bagmsg, slot))
+			arr2 = ros_numpy.numpify(msg)
+			np.testing.assert_equal(arr, arr2)
+
+	def test_pointcloud2_from_bag(self):
+		bag = Bag('test.bag')
+		# There is only one message
+		for topic, bagmsg, time in bag.read_messages(topics=['/poincloud']):
+			arr = ros_numpy.numpify(msg)
+			msg = PointCloud2()
+			for slot in msg.__slots__:
+				# Go thru all the slots, data, header...
+				setattr(msg, slot, getattr(bagmsg, slot))
+			arr2 = ros_numpy.numpify(msg)
+			np.testing.assert_equal(arr, arr2)
+
+	def test_occupancy_grid_from_bag(self):
+		bag = Bag('test.bag')
+		# There is only one message
+		for topic, bagmsg, time in bag.read_messages(topics=['/occupancygrid']):
+			arr = ros_numpy.numpify(bagmsg)
+			msg = OccupancyGrid()
+			for slot in msg.__slots__:
+				# Go thru all the slots, data, header...
+				setattr(msg, slot, getattr(bagmsg, slot))
+			arr2 = ros_numpy.numpify(msg)
+			np.testing.assert_equal(arr, arr2)
+
+	def test_vector3_from_bag(self):
+		bag = Bag('test.bag')
+		# There is only one message
+		for topic, bagmsg, time in bag.read_messages(topics=['/vector3']):
+			arr = ros_numpy.numpify(bagmsg)
+			msg = Vector3()
+			for slot in msg.__slots__:
+				# Go thru all the slots, data, header...
+				setattr(msg, slot, getattr(bagmsg, slot))
+			arr2 = ros_numpy.numpify(msg)
+			np.testing.assert_equal(arr, arr2)
+
+	def test_point_from_bag(self):
+		bag = Bag('test.bag')
+		# There is only one message
+		for topic, bagmsg, time in bag.read_messages(topics=['/point']):
+			arr = ros_numpy.numpify(bagmsg)
+			msg = Point()
+			for slot in msg.__slots__:
+				# Go thru all the slots, data, header...
+				setattr(msg, slot, getattr(bagmsg, slot))
+			arr2 = ros_numpy.numpify(msg)
+			np.testing.assert_equal(arr, arr2)
+
+	def test_quaternion_from_bag(self):
+		bag = Bag('test.bag')
+		# There is only one message
+		for topic, bagmsg, time in bag.read_messages(topics=['/quaternion']):
+			arr = ros_numpy.numpify(bagmsg)
+			msg = Quaternion()
+			for slot in msg.__slots__:
+				# Go thru all the slots, data, header...
+				setattr(msg, slot, getattr(bagmsg, slot))
+			arr2 = ros_numpy.numpify(msg)
+			np.testing.assert_equal(arr, arr2)
+
+	def test_transform_from_bag(self):
+		bag = Bag('test.bag')
+		# There is only one message
+		for topic, bagmsg, time in bag.read_messages(topics=['/transform']):
+			arr = ros_numpy.numpify(bagmsg)
+			msg = Transform()
+			for slot in msg.__slots__:
+				# Go thru all the slots, data, header...
+				setattr(msg, slot, getattr(bagmsg, slot))
+			arr2 = ros_numpy.numpify(msg)
+			np.testing.assert_equal(arr, arr2)
+
+	def test_pose_from_bag(self):
+		bag = Bag('test.bag')
+		# There is only one message
+		for topic, bagmsg, time in bag.read_messages(topics=['/pose']):
+			arr = ros_numpy.numpify(bagmsg)
+			msg = Pose()
+			for slot in msg.__slots__:
+				# Go thru all the slots, data, header...
+				setattr(msg, slot, getattr(bagmsg, slot))
+			arr2 = ros_numpy.numpify(msg)
+			np.testing.assert_equal(arr, arr2)
+
+	def test_compressed_image_from_bag(self):
+		bag = Bag('test.bag')
+		# There is only one message
+		for topic, bagmsg, time in bag.read_messages(topics=['/compressedimage']):
+			arr = ros_numpy.numpify(bagmsg)
+			msg = CompressedImage()
+			for slot in msg.__slots__:
+				# Go thru all the slots, data, header...
+				setattr(msg, slot, getattr(bagmsg, slot))
+			arr2 = ros_numpy.numpify(msg)
+			np.testing.assert_equal(arr, arr2)
+
+	def test_pose_from_bag_with_different_md5(self):
+		bag = Bag('test.bag')
+		# There is only one message
+		for topic, bagmsg, time in bag.read_messages(topics=['/pose']):
+			# use a different md5 than the installed one, for reference
+			# ros-indigo-geometry-msgs 1.11.9-0trusty-20160627-170848-0700
+			# Pose msg has md5sum e45d45a5a1ce597b249e23fb30fc871f
+			# so we use one made up from:
+			from different_md5_Pose import Pose as different_md5_Pose
+			msg = different_md5_Pose()
+			with self.assertRaises(ValueError):
+				arr = ros_numpy.numpify(msg)
+
+
+
+if __name__ == '__main__':
+	unittest.main()