diff --git a/README.md b/README.md index 35d452c..8016c06 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. 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..4f7069b 100644 --- a/src/ros_numpy/registry.py +++ b/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( 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/test.bag b/test/test.bag new file mode 100644 index 0000000..9c2d79a 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..1551ff3 --- /dev/null +++ b/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() 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()