(tf2_pykdl_ros) add more convert functions

tf2_pykdl_ros/src/tf2_pykdl_ros/__init__.py

@@ -47,24 +47,27 @@ def to_msg_vector(vector: VectorStamped) -> PointStamped:
 tf2_ros.ConvertRegistration().add_to_msg(VectorStamped, to_msg_vector)
 
 
-def from_msg_vector(msg: Union[PointStamped, PoseStamped]) -> VectorStamped:
+def from_msg_vector(msg: Union[PointStamped, PoseStamped, TransformStamped]) -> VectorStamped:
     """
-    Convert a PointStamped/PoseStamped message to a stamped VectorStamped.
+    Convert a PointStamped/PoseStamped/TransformStamped message to a stamped VectorStamped.
 
-    :param msg: The PointStamped/PoseStamped message to convert.
+    :param msg: The PointStamped/PoseStamped/TransformStamped message to convert.
     :return: The timestamped converted PyKDL vector.
     """
     if isinstance(msg, PointStamped):
         v = msg.point
     elif isinstance(msg, PoseStamped):
         v = msg.pose.position
+    elif isinstance(msg, TransformStamped):
+        v = msg.transform.translation
     else:
-        raise TypeError(f"msg should be PointStamped or PoseStamped, not '{type(msg)}'")
+        raise TypeError(f"msg should be PointStamped, PoseStamped or TransformStamped, not '{type(msg)}'")
     return VectorStamped.from_xyz(v.x, v.y, v.z, msg.header.stamp, msg.header.frame_id)
 
 
 tf2_ros.ConvertRegistration().add_convert((PointStamped, VectorStamped), from_msg_vector)
 tf2_ros.ConvertRegistration().add_convert((PoseStamped, VectorStamped), from_msg_vector)
+tf2_ros.ConvertRegistration().add_convert((TransformStamped, VectorStamped), from_msg_vector)
 tf2_ros.ConvertRegistration().add_from_msg(VectorStamped, from_msg_vector)
 
 
@@ -97,7 +100,26 @@ def do_transform_vector(vector: VectorStamped, transform: TransformStamped) -> V
 tf2_ros.TransformRegistration().add(VectorStamped, do_transform_vector)
 
 
-def to_msg_frame(frame: FrameStamped) -> PoseStamped:
+def to_point_msg_frame(frame: FrameStamped) -> PointStamped:
+    """
+    Convert a FrameStamped to a geometry_msgs PointStamped message.
+
+    :param frame: The frame to convert.
+    :return: The converted point.
+    """
+    msg = PointStamped()
+    msg.header = frame.header
+    v = frame.frame.p
+    msg.point.x = v[0]
+    msg.point.y = v[1]
+    msg.point.z = v[2]
+    return msg
+
+
+tf2_ros.ConvertRegistration().add_convert((FrameStamped, PointStamped), to_point_msg_frame)
+
+
+def to_pose_msg_frame(frame: FrameStamped) -> PoseStamped:
     """
     Convert a FrameStamped to a geometry_msgs PoseStamped message.
 
@@ -118,28 +140,76 @@ def to_msg_frame(frame: FrameStamped) -> PoseStamped:
     return msg
 
 
-tf2_ros.ConvertRegistration().add_convert((FrameStamped, PoseStamped), to_msg_frame)
-tf2_ros.ConvertRegistration().add_to_msg(FrameStamped, to_msg_frame)
+tf2_ros.ConvertRegistration().add_convert((FrameStamped, PoseStamped), to_pose_msg_frame)
+tf2_ros.ConvertRegistration().add_to_msg(FrameStamped, to_pose_msg_frame)
+
+
+def to_transform_msg_frame(frame: FrameStamped) -> TransformStamped:
+    """
+    Convert a FrameStamped to a geometry_msgs TransformStamped message.
+
+    The child_frame_id is not set.
+
+    :param frame: The frame to convert.
+    :return: The converted Pose.
+    """
+    msg = TransformStamped()
+    msg.header = frame.header
+    v = frame.frame.p
+    msg.transform.translation.x = v[0]
+    msg.transform.translation.y = v[1]
+    msg.transform.translation.z = v[2]
+    q = frame.frame.M.GetQuaternion()
+    msg.transform.rotation.x = q[0]
+    msg.transform.rotation.y = q[1]
+    msg.transform.rotation.z = q[2]
+    msg.transform.rotation.w = q[3]
+    return msg
+
+
+tf2_ros.ConvertRegistration().add_convert((FrameStamped, TransformStamped), to_transform_msg_frame)
 
 
-def from_msg_frame(msg: PoseStamped) -> FrameStamped:
+def from_msg_frame(msg: Union[PointStamped, PoseStamped, TransformStamped]) -> FrameStamped:
     """
     Convert a PoseStamped message to a stamped FrameStamped.
 
     :param msg: The PoseStamped message to convert.
     :return: The timestamped converted PyKDL frame.
     """
-    if not isinstance(msg, PoseStamped):
-        raise TypeError(f"msg should be PoseStamped, not '{type(msg)}'")
-    v = msg.pose.position
-    vector = kdl.Vector(v.x, v.y, v.z)
-    q = msg.pose.orientation
-    rotation = kdl.Rotation.Quaternion(q.x, q.y, q.z, q.w)
+    if (
+        not isinstance(msg, PointStamped)
+        and not isinstance(msg, PoseStamped)
+        and not isinstance(msg, TransformStamped)
+    ):
+        raise TypeError(f"msg should be PointStamped, PoseStamped or TransformStamped, not '{type(msg)}'")
+    if isinstance(msg, PointStamped):
+        pose_attr = "point"
+    elif isinstance(msg, PoseStamped):
+        pose_attr = "pose"
+        position_attr = "position"
+        orientation_attr = "orientation"
+    elif isinstance(msg, TransformStamped):
+        pose_attr = "transform"
+        position_attr = "translation"
+        orientation_attr = "rotation"
+
+    pose = getattr(msg, pose_attr)
+    if isinstance(msg, PointStamped):
+        vector = kdl.Vector(pose.x, pose.y, pose.z)
+        rotation = kdl.Rotation.Identity()
+    else:
+        v = getattr(pose, position_attr)
+        vector = kdl.Vector(v.x, v.y, v.z)
+        q = getattr(pose, orientation_attr)
+        rotation = kdl.Rotation.Quaternion(q.x, q.y, q.z, q.w)
     frame = kdl.Frame(rotation, vector)
     return FrameStamped(frame, msg.header.stamp, msg.header.frame_id)
 
 
+tf2_ros.ConvertRegistration().add_convert((PointStamped, FrameStamped), from_msg_frame)
 tf2_ros.ConvertRegistration().add_convert((PoseStamped, FrameStamped), from_msg_frame)
+tf2_ros.ConvertRegistration().add_convert((TransformStamped, FrameStamped), from_msg_frame)
 tf2_ros.ConvertRegistration().add_from_msg(FrameStamped, from_msg_frame)