Support CompressedImage and Stamped types

README.md

@@ -26,12 +26,14 @@ 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`.
 
 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):

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