# -*- coding: UTF8 -*-
"""
TF topic handling
author: Michael Grupp
This file is part of evo (github.com/MichaelGrupp/evo).
evo is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
evo is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with evo. If not, see .
"""
import dataclasses
import logging
import math
import warnings
from collections import defaultdict
from typing import (
DefaultDict,
Protocol,
runtime_checkable,
)
import numpy as np
import tf2_py
from geometry_msgs.msg import TransformStamped
from rosbags.rosbag1 import Reader as Rosbag1Reader
from rosbags.rosbag2 import Reader as Rosbag2Reader
from rosbags.typesys import get_typestore, get_types_from_msg, Stores
from rosbags.typesys.store import Typestore
from evo import EvoException
from evo.core.trajectory import PoseTrajectory3D
from evo.tools import tf_id
from evo.tools.file_interface import _get_xyz_quat_from_transform_stamped
from evo.tools.settings import SETTINGS
logger = logging.getLogger(__name__)
SUPPORTED_TF_MSG = "tf2_msgs/msg/TFMessage"
DEFAULT_MAX_TIME = SETTINGS.tf_cache_max_time
DEFAULT_DEBUG = SETTINGS.tf_cache_debug
class TfCacheException(EvoException):
pass
@runtime_checkable
class Ros1TimeLike(Protocol): # pylint: disable=too-few-public-methods
"""
Basic ROS 1 compatible time instance protocol.
"""
def to_sec(self) -> float:
"""
Gets scalar time, in seconds.
"""
def __lt__(self, other) -> bool:
"""less-than comparator, required for sorting"""
@runtime_checkable
class Ros2TimeLike(Protocol): # pylint: disable=too-few-public-methods
"""
Basic ROS 2 compatible time instance protocol.
"""
@property
def nanoseconds(self) -> int:
"""
Gets underlying scalar timestamp, in nanoseconds.
"""
def __lt__(self, other) -> bool:
"""less-than comparator, required for sorting"""
@runtime_checkable
class Ros2StampLike(Protocol): # pylint: disable=too-few-public-methods
"""
Basic ROS 2 compatible message stamp protocol.
"""
sec: int
nanosec: int
@dataclasses.dataclass
class TfDuration(Ros1TimeLike, Ros2TimeLike, Ros2StampLike):
"""
A duration representation that is TF compatible in ROS 1 and ROS 2.
"""
sec: int
nanosec: int
@classmethod
def from_sec(cls, sec: float) -> "TfDuration":
"""Instantiates a duration given a number of seconds."""
frac, whole = math.modf(sec)
return cls(sec=int(whole), nanosec=int(frac * 1e9))
@property
def nanoseconds(self) -> int:
return self.sec * 1000000000 + self.nanosec
def to_sec(self) -> float:
return self.sec + self.nanosec * 1e-9
def to_sec(
timestamp: Ros1TimeLike | Ros2TimeLike | Ros2StampLike,
) -> float:
"""Converts any given `timestamp` to a scalar time, in seconds."""
if isinstance(timestamp, Ros1TimeLike):
return timestamp.to_sec()
if isinstance(timestamp, Ros2TimeLike):
return timestamp.nanoseconds * 1e-9
return timestamp.sec + timestamp.nanosec * 1e-9
@dataclasses.dataclass
class TimeSpec:
"""
Internal timespec for time tracking.
"""
sec: int
nanosec: int
@dataclasses.dataclass
class CacheKey:
"""
Key for presence tracking in TfCache.
"""
reader: Rosbag1Reader | Rosbag2Reader
topic: str
def __hash__(self):
"""Hash for a bag & TF topic pair."""
return hash(f"{self.reader.path}+{self.topic}")
class TfCache(object):
"""
For caching TF messages and looking up trajectories of specific transforms.
Allows to use multiple bag files and TF topics.
"""
def __init__(
self,
*,
max_time_sec: float = DEFAULT_MAX_TIME,
debug: bool = DEFAULT_DEBUG,
) -> None:
"""
:param max_time_sec: Maximum TF buffer size in seconds.
:param debug: Whether to emit debug logs with TF buffer details.
"""
self.buffer = tf2_py.BufferCore(TfDuration.from_sec(max_time_sec))
# Cache presence is tracked by storing the start time of each TF topic/bag combo.
self.cache: dict[CacheKey, TimeSpec] = {}
self.debug = debug and logger.isEnabledFor(logging.DEBUG)
def clear(self) -> None:
logger.debug("Clearing TF cache.")
self.buffer.clear()
self.cache = {}
# tf2_msgs/TFMessage is not included in default rosbags typestore,
# update the ROS1 typestore with the interface definition from the bag.
# https://ternaris.gitlab.io/rosbags/examples/register_types.html
@staticmethod
def _setup_typestore(
reader: Rosbag1Reader | Rosbag2Reader,
) -> Typestore:
if isinstance(reader, Rosbag2Reader):
return get_typestore(Stores.LATEST)
typestore = get_typestore(Stores.ROS1_NOETIC)
for connection in reader.connections:
if connection.msgtype == SUPPORTED_TF_MSG:
typestore.register(
get_types_from_msg(
connection.msgdef.data, connection.msgtype
)
)
break
return typestore
def from_bag(
self,
reader: Rosbag1Reader | Rosbag2Reader,
topic: str = "/tf",
static_topic: str = "/tf_static",
) -> None:
"""
Loads the TF topics from a bagfile into the buffer,
if it's not already cached.
:param reader: opened bag reader (rosbags.rosbag1)
:param topic: TF topic
"""
tf_topics = [topic]
if topic not in reader.topics:
raise TfCacheException(f"no messages for topic {topic} in bag")
# Implicitly add static TFs to buffer if present.
if static_topic in reader.topics:
tf_topics.append(static_topic)
typestore = self._setup_typestore(reader)
# Add TF data to buffer if this bag/topic pair is not already cached.
for tf_topic in tf_topics:
cache_key = CacheKey(reader, tf_topic)
if cache_key in self.cache:
logger.debug(
f"Using cache for topic {tf_topic} from {reader.path}"
)
continue
logger.debug(f"Caching TF topic {tf_topic} from {reader.path} ...")
start_time = None
connections = [
c for c in reader.connections if c.topic == tf_topic
]
for connection, _, rawdata in reader.messages(
connections=connections
):
if connection.msgtype != SUPPORTED_TF_MSG:
raise TfCacheException(
f"Expected {SUPPORTED_TF_MSG} message type for topic "
f"{tf_topic}, got: {connection.msgtype}"
)
if isinstance(reader, Rosbag1Reader):
msg = typestore.deserialize_ros1(
rawdata, connection.msgtype
)
else:
msg = typestore.deserialize_cdr(
rawdata, connection.msgtype
)
for tf in msg.transforms: # type: ignore
# Store the time of the first TF message.
# This is more robust than reader.start_time,
# e.g. when sim time was used (see also: #771)
if start_time is None:
start_time = TimeSpec(
tf.header.stamp.sec, tf.header.stamp.nanosec
)
self.cache[cache_key] = start_time
# Convert from rosbags.typesys.types to native ROS.
# Related: https://gitlab.com/ternaris/rosbags/-/issues/13
native_msg = TransformStamped()
if hasattr(native_msg.header.stamp, "nsecs"):
native_msg.header.stamp.secs = tf.header.stamp.sec
native_msg.header.stamp.nsecs = tf.header.stamp.nanosec
else:
native_msg.header.stamp.sec = tf.header.stamp.sec
native_msg.header.stamp.nanosec = (
tf.header.stamp.nanosec
)
native_msg.header.frame_id = tf.header.frame_id
native_msg.child_frame_id = tf.child_frame_id
native_msg.transform.translation.x = (
tf.transform.translation.x
)
native_msg.transform.translation.y = (
tf.transform.translation.y
)
native_msg.transform.translation.z = (
tf.transform.translation.z
)
native_msg.transform.rotation.x = tf.transform.rotation.x
native_msg.transform.rotation.y = tf.transform.rotation.y
native_msg.transform.rotation.z = tf.transform.rotation.z
native_msg.transform.rotation.w = tf.transform.rotation.w
if tf_topic == static_topic:
self.buffer.set_transform_static(native_msg, __name__)
else:
self.buffer.set_transform(native_msg, __name__)
if self.debug:
logger.debug("TF buffer: %s", self.buffer.all_frames_as_yaml())
def lookup_trajectory(
self,
parent_frame: str,
child_frame: str,
timestamps: list[Ros1TimeLike] | list[Ros2TimeLike],
) -> PoseTrajectory3D:
"""
Look up the trajectory of a transform chain from the cache's TF buffer.
:param parent_frame, child_frame: TF transform frame IDs
:param timestamps: timestamps at which to lookup the trajectory poses.
"""
stamps, xyz, quat = [], [], []
timestamps.sort()
for timestamp in timestamps:
try:
tf = self.buffer.lookup_transform_core(
parent_frame, child_frame, timestamp
)
except tf2_py.ExtrapolationException:
continue
stamps.append(to_sec(tf.header.stamp))
x, q = _get_xyz_quat_from_transform_stamped(tf)
xyz.append(x)
quat.append(q)
trajectory = PoseTrajectory3D(
np.array(xyz),
np.array(quat),
np.array(stamps),
meta={"frame_id": parent_frame, "child_frame_id": child_frame},
)
return trajectory
def get_trajectory(
self,
reader: Rosbag1Reader | Rosbag2Reader,
identifier: str,
timestamps: list[Ros1TimeLike] | list[Ros2TimeLike] | None = None,
) -> PoseTrajectory3D:
"""
Get a TF trajectory from a bag file. Updates or uses the cache.
:param reader: opened bag reader (rosbags.rosbag1)
:param identifier: trajectory ID :.
Example: /tf:map.base_link
"""
split_id = tf_id.split_id(identifier)
topic, parent, child = split_id[0], split_id[1], split_id[2]
static_topic = split_id[3] if len(split_id) == 4 else "/tf_static"
logger.debug(
f"Loading trajectory of transform '{parent} to {child}' "
f"from topic {topic} (static topic: {static_topic})."
)
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=UserWarning)
self.from_bag(reader, topic, static_topic)
if timestamps is None:
timestamps = []
try:
latest_time = self.buffer.get_latest_common_time(parent, child)
except (
tf2_py.LookupException,
tf2_py.TransformException,
) as tf_error:
error = f"Could not load trajectory: {tf_error}"
if not self.debug:
error += (
"\nEnable 'evo_config set tf_cache_debug true' and "
"verbose/debug logging to see TF buffer contents."
)
raise TfCacheException(error) from tf_error
_start_time: TimeSpec = self.cache[CacheKey(reader, topic)]
if hasattr(latest_time, "nsecs"):
from rospy import Time, Duration
start_time = Time(_start_time.sec, _start_time.nanosec)
step = Duration.from_sec(
1.0 / SETTINGS.tf_cache_lookup_frequency
)
else:
from rclpy.time import Time
from rclpy.duration import Duration
start_time = Time(
seconds=_start_time.sec, nanoseconds=_start_time.nanosec
)
step = Duration(
seconds=1.0 / SETTINGS.tf_cache_lookup_frequency
)
# Static TF have zero timestamp in the buffer, which will be lower
# than the bag start time. Looking up a static TF is a valid request,
# so this should be possible.
if latest_time < start_time:
timestamps.append(latest_time)
else:
time = start_time
while time <= latest_time:
timestamps.append(time)
time = time + step
return self.lookup_trajectory(parent, child, timestamps)
__instance: DefaultDict[int, TfCache] = defaultdict(lambda: TfCache())
def instance(hash: int) -> TfCache:
"""Hacky module-level "singleton" of TfCache"""
global __instance
return __instance[hash]