diff --git a/README.md b/README.md index 35d452c..e6e99a5 100644 --- a/README.md +++ b/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: diff --git a/src/ros_numpy/geometry.py b/src/ros_numpy/geometry.py index bd9a4ec..83f3c95 100644 --- a/src/ros_numpy/geometry.py +++ b/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 diff --git a/src/ros_numpy/image.py b/src/ros_numpy/image.py index bfd3a1b..a4895ef 100644 --- a/src/ros_numpy/image.py +++ b/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): diff --git a/src/ros_numpy/registry.py b/src/ros_numpy/registry.py index 5442b93..a0a7fb7 100644 --- a/src/ros_numpy/registry.py +++ b/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( diff --git a/test/__init__.py b/test/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test/create_test_bag.py b/test/create_test_bag.py new file mode 100644 index 0000000..97e14bf --- /dev/null +++ b/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() diff --git a/test/different_md5_Pose.py b/test/different_md5_Pose.py new file mode 100644 index 0000000..d02205f --- /dev/null +++ b/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") diff --git a/test/test.bag b/test/test.bag new file mode 100644 index 0000000..254dbc3 Binary files /dev/null and b/test/test.bag differ diff --git a/test/test_from_bag.py b/test/test_from_bag.py new file mode 100644 index 0000000..b7fcf8f --- /dev/null +++ b/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() diff --git a/test/test_geometry.py b/test/test_geometry.py index 12aa682..422197e 100644 --- a/test/test_geometry.py +++ b/test/test_geometry.py @@ -5,6 +5,8 @@ from tf import transformations from geometry_msgs.msg import Vector3, Quaternion, Transform, Point, Pose +from geometry_msgs.msg import TransformStamped, Vector3Stamped, QuaternionStamped, PointStamped, PoseStamped + class TestGeometry(unittest.TestCase): def test_point(self): @@ -75,5 +77,81 @@ def test_pose(self): np.testing.assert_allclose(msg.orientation.z, t.orientation.z) np.testing.assert_allclose(msg.orientation.w, t.orientation.w) + def test_pointstamped(self): + ps = PointStamped() + p = Point(1, 2, 3) + ps.point = p + + p_arr = ros_numpy.numpify(ps) + np.testing.assert_array_equal(p_arr, [1, 2, 3]) + + p_arrh = ros_numpy.numpify(ps, hom=True) + np.testing.assert_array_equal(p_arrh, [1, 2, 3, 1]) + + self.assertEqual(p, ros_numpy.msgify(Point, p_arr)) + self.assertEqual(p, ros_numpy.msgify(Point, p_arrh)) + self.assertEqual(p, ros_numpy.msgify(Point, p_arrh * 2)) + + def test_vector3stamped(self): + vs = Vector3Stamped() + v = Vector3(1, 2, 3) + vs.vector = v + + v_arr = ros_numpy.numpify(vs) + np.testing.assert_array_equal(v_arr, [1, 2, 3]) + + v_arrh = ros_numpy.numpify(vs, hom=True) + np.testing.assert_array_equal(v_arrh, [1, 2, 3, 0]) + + self.assertEqual(v, ros_numpy.msgify(Vector3, v_arr)) + self.assertEqual(v, ros_numpy.msgify(Vector3, v_arrh)) + + with self.assertRaises(AssertionError): + ros_numpy.msgify(Vector3, np.array([0, 0, 0, 1])) + + def test_transformstamped(self): + ts = TransformStamped() + t = Transform( + translation=Vector3(1, 2, 3), + rotation=Quaternion(*transformations.quaternion_from_euler(np.pi, 0, 0)) + ) + ts.transform = t + + t_mat = ros_numpy.numpify(ts) + + np.testing.assert_allclose(t_mat.dot([0, 0, 1, 1]), [1.0, 2.0, 2.0, 1.0]) + + msg = ros_numpy.msgify(Transform, t_mat) + + np.testing.assert_allclose(msg.translation.x, t.translation.x) + np.testing.assert_allclose(msg.translation.y, t.translation.y) + np.testing.assert_allclose(msg.translation.z, t.translation.z) + np.testing.assert_allclose(msg.rotation.x, t.rotation.x) + np.testing.assert_allclose(msg.rotation.y, t.rotation.y) + np.testing.assert_allclose(msg.rotation.z, t.rotation.z) + np.testing.assert_allclose(msg.rotation.w, t.rotation.w) + + def test_posestamped(self): + pss = PoseStamped() + ps = Pose( + position=Point(1.0, 2.0, 3.0), + orientation=Quaternion(*transformations.quaternion_from_euler(np.pi, 0, 0)) + ) + pss.pose = ps + + t_mat = ros_numpy.numpify(pss) + + np.testing.assert_allclose(t_mat.dot([0, 0, 1, 1]), [1.0, 2.0, 2.0, 1.0]) + + msg = ros_numpy.msgify(Pose, t_mat) + + np.testing.assert_allclose(msg.position.x, ps.position.x) + np.testing.assert_allclose(msg.position.y, ps.position.y) + np.testing.assert_allclose(msg.position.z, ps.position.z) + np.testing.assert_allclose(msg.orientation.x, ps.orientation.x) + np.testing.assert_allclose(msg.orientation.y, ps.orientation.y) + np.testing.assert_allclose(msg.orientation.z, ps.orientation.z) + np.testing.assert_allclose(msg.orientation.w, ps.orientation.w) + if __name__ == '__main__': unittest.main()