Add rosbag support, add CompressedImage, add geometry_msgs Stamped types

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.
 
 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

@@ -2,6 +2,7 @@
 import genpy
 import collections
 from . import numpy_msg
+from importlib import import_module
 
 _to_numpy = {}
 _from_numpy = {}
@@ -32,6 +33,15 @@ 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:
+		# a rosbag message class name looks like '_sensor_msgs__Image'
+		module_name, msg_class_name = msg.__class__.__name__.split('__')
+		if module_name and msg_class_name:
+			convertable_class_names = [c.__name__ for c in zip(*_to_numpy.keys())[0]]
+			module_name = module_name.lstrip('_') 
+			module = import_module(module_name + '.msg')
+			class_ = getattr(module, msg_class_name)
+			conv = _to_numpy.get((class_, False))
 
 	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/test_from_bag.py

@@ -0,0 +1,122 @@
+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)
+
+
+if __name__ == '__main__':
+	unittest.main()