diff --git a/src/imetro_behavior/imetro_behavior/geometry_behaviors.py b/src/imetro_behavior/imetro_behavior/geometry_behaviors.py index 1790804..c77e032 100644 --- a/src/imetro_behavior/imetro_behavior/geometry_behaviors.py +++ b/src/imetro_behavior/imetro_behavior/geometry_behaviors.py @@ -762,3 +762,49 @@ def terminate(self, new_status: Status) -> None: """Cleanup the publisher""" if self.status == Status.RUNNING and new_status == Status.INVALID and self.publisher is not None: self.publisher.destroy() + + +class InvertTransformStamped(BehaviourWithPorts): + """Invert the TransformStamped message.""" + + INPUT_PORTS = { + "transform_stamped_in": PortInformation(data_type=TransformStamped, required=True), + } + + OUTPUT_PORTS = {"transform_stamped_out": PortInformation(data_type=TransformStamped, required=True)} + + def update(self) -> Status: + """Convert TransformStamped message into a 4x4 matrix, invert the matrix, and convert the result back to TransformStamped message.""" + tf_in = self.get_input("transform_stamped_in") + + xyz_in = [tf_in.transform.translation.x, tf_in.transform.translation.y, tf_in.transform.translation.z] + xyzw_in = [ + tf_in.transform.rotation.x, + tf_in.transform.rotation.y, + tf_in.transform.rotation.z, + tf_in.transform.rotation.w, + ] + + T_in = np.eye(4) + T_in[:3, :3] = R.from_quat(xyzw_in).as_matrix() + T_in[:3, 3] = xyz_in + T_out = np.linalg.inv(T_in) + + xyz_out = T_out[:3, 3] + xyzw_out = R.from_matrix(T_out[:3, :3]).as_quat() + + tf_out = TransformStamped() + tf_out.header.stamp = tf_in.header.stamp + tf_out.header.frame_id = tf_in.child_frame_id + tf_out.child_frame_id = tf_in.header.frame_id + tf_out.transform.translation.x = xyz_out[0] + tf_out.transform.translation.y = xyz_out[1] + tf_out.transform.translation.z = xyz_out[2] + + tf_out.transform.rotation.x = xyzw_out[0] + tf_out.transform.rotation.y = xyzw_out[1] + tf_out.transform.rotation.z = xyzw_out[2] + tf_out.transform.rotation.w = xyzw_out[3] + + self._set_output("transform_stamped_out", tf_out) + return Status.SUCCESS diff --git a/src/imetro_behavior/tests/test_geometry_behaviors.py b/src/imetro_behavior/tests/test_geometry_behaviors.py index 392d26e..460cf48 100644 --- a/src/imetro_behavior/tests/test_geometry_behaviors.py +++ b/src/imetro_behavior/tests/test_geometry_behaviors.py @@ -32,6 +32,7 @@ DecomposePoseStamped, GetRelativePoseStamped, GetRollPitchYaw, + InvertTransformStamped, LookupTransform, OffsetPoseStamped, PoseStampedToTransformStamped, @@ -650,3 +651,50 @@ def test_decompose_pose_stamped() -> None: assert frame_id_msg == "odom" assert translation_msg == [1.0, 2.0, 3.0] assert orientation_msg == [0.0, 0.0, 0.0, 1.0] + + +def test_invert_transport_stamped() -> None: + behavior = InvertTransformStamped(name="invert_transform_stamped") + behavior.setup_ports() + + for _ in range(10): + a_T_b = np.eye(4) + a_T_b[:3, :3] = R.random().as_matrix() + a_T_b[:3, 3] = np.random.rand( + 3, + ) + + a_t_b = a_T_b[:3, 3] + a_R_b = R.from_matrix(a_T_b[:3, :3]).as_quat() + + tf_in = TransformStamped() + tf_in.header.frame_id = "frame_a" + tf_in.child_frame_id = "frame_b" + tf_in.transform.translation.x = a_t_b[0] + tf_in.transform.translation.y = a_t_b[1] + tf_in.transform.translation.z = a_t_b[2] + + tf_in.transform.rotation.x = a_R_b[0] + tf_in.transform.rotation.y = a_R_b[1] + tf_in.transform.rotation.z = a_R_b[2] + tf_in.transform.rotation.w = a_R_b[3] + + Blackboard.set(behavior._get_blackboard_key("transform_stamped_in"), tf_in) + behavior.tick_once() + assert behavior.status == Status.SUCCESS + tf_out = behavior.get_last_output("transform_stamped_out") + assert tf_out.header.frame_id == tf_in.child_frame_id + assert tf_out.child_frame_id == tf_in.header.frame_id + + b_T_a = np.linalg.inv(a_T_b) + b_t_a = b_T_a[:3, 3] + b_R_a = R.from_matrix(b_T_a[:3, :3]).as_quat() + + assert tf_out.transform.translation.x == pytest.approx(b_t_a[0], abs=1e-5) + assert tf_out.transform.translation.y == pytest.approx(b_t_a[1], abs=1e-5) + assert tf_out.transform.translation.z == pytest.approx(b_t_a[2], abs=1e-5) + + assert tf_out.transform.rotation.x == pytest.approx(b_R_a[0], abs=1e-5) + assert tf_out.transform.rotation.y == pytest.approx(b_R_a[1], abs=1e-5) + assert tf_out.transform.rotation.z == pytest.approx(b_R_a[2], abs=1e-5) + assert tf_out.transform.rotation.w == pytest.approx(b_R_a[3], abs=1e-5)