feat: Better tf_static insert

src/kappe/convert.py

@@ -17,7 +17,7 @@
 
 from kappe import __version__
 from kappe.module.pointcloud import point_cloud
-from kappe.module.tf import tf_remove, tf_static_insert
+from kappe.module.tf import TF_SCHEMA_NAME, TF_SCHEMA_TEXT, tf_remove, tf_static_insert
 from kappe.module.timing import fix_ros1_time, time_offset
 from kappe.plugin import ConverterPlugin, load_plugin
 from kappe.settings import Settings
@@ -96,6 +96,8 @@ def __init__(self, config: Settings, input_path: Path, output_path: Path, raw_co
 
         self.statistics: Statistics = self.summary.statistics
 
+        self.tf_static_channel_id: int | None = None
+
         # mapping of schema name to schema
         self.schema_list: dict[str, Schema] = {}
         self.init_schema()
@@ -153,6 +155,12 @@ def init_schema(self):
                 self.schema_list[out_schema] = self.writer.register_msgdef(
                     out_schema, new_data)
 
+        if self.config.tf_static and TF_SCHEMA_NAME not in self.schema_list:
+            # insert tf schema
+            self.schema_list[TF_SCHEMA_NAME] = self.writer.register_msgdef(
+                TF_SCHEMA_NAME, TF_SCHEMA_TEXT,
+            )
+
     def init_channel(self):
         for channel in self.summary.channels.values():
             metadata = channel.metadata
@@ -201,6 +209,14 @@ def init_channel(self):
                 )
                 self.writer._channel_ids[topic] = channel_id  # noqa: SLF001
 
+        if self.config.tf_static and '/tf_static' not in self.writer._channel_ids:  # noqa: SLF001
+            # insert tf schema
+            channel_id = self.writer._writer.register_channel(  # noqa: SLF001
+                topic='/tf_static',
+                message_encoding='cdr',
+                schema_id=self.schema_list[TF_SCHEMA_NAME].id,
+            )
+
     def get_selected_channels(self) -> set[str]:
         """Get a list of channels that should be converted."""
         filtered_channels: set[str] = set()
@@ -301,9 +317,6 @@ def process_message(self, msg: DecodedMessageTuple):
         if topic in ['/tf', '/tf_static']:
             tf_remove(self.config.tf_static, msg)
 
-        if topic in ['/tf_static'] and not self.tf_inserted:
-            self.tf_inserted = tf_static_insert(self.config.tf_static, msg)
-
         if topic in self.config.time_offset:
             time_offset(self.config.time_offset[topic], msg)
 
@@ -328,6 +341,10 @@ def process_file(self, tqdm_idx: int = 0):
         if self.config.time_start is not None:
             start_time = max(start_time, self.config.time_start)
 
+        start_time_part_sec = int(start_time)
+        start_time_part_ns = int((start_time - start_time_part_sec) * 1e9)
+        start_time_ns = int(start_time * 1e9)
+
         end_time = self.statistics.message_end_time / 1e9
         if self.config.time_end is not None:
             conf_end_time = self.config.time_end
@@ -354,27 +371,26 @@ def process_file(self, tqdm_idx: int = 0):
             if tf_static_iter is None:
                 raise ValueError('tf_static_iter is None')
 
-            secs = int(start_time)
-            nsecs = int((start_time - secs) * 1e9)
-            nano_int = int(start_time * 1e9)
-
             for count, msg in enumerate(tf_static_iter, 1):
+                if msg.schema is None:
+                    continue
                 # patch header stamp
-                for transform in msg.ros_msg.transforms:
+                ros_msg = msg.decoded_message
+                for transform in ros_msg.transforms:
                     # foxglove does not tf msg with the exact same timestamp
-                    nsecs += 1
-                    nano_int += 1
+                    start_time_part_ns += 1
+                    start_time_ns += 1
 
-                    transform.header.stamp.sec = secs
-                    transform.header.stamp.nanosec = nsecs
+                    transform.header.stamp.sec = start_time_part_sec
+                    transform.header.stamp.nanosec = start_time_part_ns
 
                 self.writer.write_message(
                     topic=msg.channel.topic,
                     schema=self.schema_list[msg.schema.name],
-                    message=msg.ros_msg,
-                    log_time=nano_int,
-                    publish_time=nano_int,
-                    sequence=msg.sequence_count,
+                    message=ros_msg,
+                    log_time=start_time_ns,
+                    publish_time=start_time_ns,
+                    sequence=msg.message.sequence,
                 )
 
                 # performance hack
@@ -400,6 +416,17 @@ def process_file(self, tqdm_idx: int = 0):
 
         duration = end_time - start_time
 
+        # insert tf_static messages at the beginning of the file
+        insert_tf = tf_static_insert(self.config.tf_static, start_time_ns)
+        if insert_tf is not None:
+            self.writer.write_message(
+                topic='/tf_static',
+                schema=self.schema_list[TF_SCHEMA_NAME],
+                message=insert_tf,
+                log_time=start_time_ns,
+                publish_time=start_time_ns,
+            )
+
         with tqdm(
             total=duration,
             position=tqdm_idx,

src/kappe/module/tf.py

@@ -1,10 +1,87 @@
-import copy
+from typing import Any
 
 from mcap.reader import DecodedMessageTuple
 from pydantic import BaseModel
 
 from kappe.utils.settings import SettingRotation, SettingTranslation
-from kappe.utils.types import ClassDict
+
+TF_SCHEMA_NAME = 'tf2_msgs/msg/TFMessage'
+TF_SCHEMA_TEXT = """
+geometry_msgs/TransformStamped[] transforms
+
+================================================================================
+MSG: geometry_msgs/TransformStamped
+# This expresses a transform from coordinate frame header.frame_id
+# to the coordinate frame child_frame_id at the time of header.stamp
+#
+# This message is mostly used by the
+# <a href="https://index.ros.org/p/tf2/">tf2</a> package.
+# See its documentation for more information.
+#
+# The child_frame_id is necessary in addition to the frame_id
+# in the Header to communicate the full reference for the transform
+# in a self contained message.
+
+# The frame id in the header is used as the reference frame of this transform.
+std_msgs/Header header
+
+# The frame id of the child frame to which this transform points.
+string child_frame_id
+
+# Translation and rotation in 3-dimensions of child_frame_id from header.frame_id.
+Transform transform
+
+================================================================================
+MSG: geometry_msgs/Transform
+# This represents the transform between two coordinate frames in free space.
+
+Vector3 translation
+Quaternion rotation
+
+================================================================================
+MSG: geometry_msgs/Quaternion
+# This represents an orientation in free space in quaternion form.
+
+float64 x 0
+float64 y 0
+float64 z 0
+float64 w 1
+
+================================================================================
+MSG: geometry_msgs/Vector3
+# This represents a vector in free space.
+
+# This is semantically different than a point.
+# A vector is always anchored at the origin.
+# When a transform is applied to a vector, only the rotational component is applied.
+
+float64 x
+float64 y
+float64 z
+
+================================================================================
+MSG: std_msgs/Header
+# Standard metadata for higher-level stamped data types.
+# This is generally used to communicate timestamped data
+# in a particular coordinate frame.
+
+# Two-integer timestamp that is expressed as seconds and nanoseconds.
+builtin_interfaces/Time stamp
+
+# Transform frame with which this data is associated.
+string frame_id
+
+================================================================================
+MSG: builtin_interfaces/Time
+# This message communicates ROS Time defined here:
+# https://design.ros2.org/articles/clock_and_time.html
+
+# The seconds component, valid over all int32 values.
+int32 sec
+
+# The nanoseconds component, valid in the range [0, 10e9).
+uint32 nanosec
+"""
 
 
 class SettingTFInsert(BaseModel):
@@ -35,26 +112,16 @@ class SettingTF(BaseModel):
     insert: list[SettingTFInsert] | None = None
 
 
-def tf_static_insert(cfg: SettingTF, msg: DecodedMessageTuple) -> bool:
-    schema, channel, message, ros_msg = msg
-
+def tf_static_insert(cfg: SettingTF, stamp_ns: int) -> None | Any:
     if cfg.insert is None:
-        return True
+        return None
 
-    # One transform needed as reference
-    if len(ros_msg.transforms) == 0:
-        return False
+    sec = int(stamp_ns / 1e9)
+    nanosec = int(stamp_ns % 1e9)
 
-    header = ros_msg.transforms[0].header
+    transforms = []
 
     for insert in cfg.insert:
-        tf_msg = ClassDict(
-            header=copy.deepcopy(header),
-            child_frame_id=insert.child_frame_id,
-        )
-
-        tf_msg.header.frame_id = insert.frame_id
-
         trans = {}
         translation = insert.translation
         if translation is not None:
@@ -73,10 +140,20 @@ def tf_static_insert(cfg: SettingTF, msg: DecodedMessageTuple) -> bool:
                 'w': rot_quat[3],
             }
 
-        tf_msg['transform'] = trans
-        ros_msg.transforms.append(tf_msg)
-
-    return True
+        tf_msg = {
+            'header': {
+                'frame_id': insert.frame_id,
+                'stamp': {
+                    'sec': sec,
+                    'nanosec': nanosec,
+                },
+            },
+            'child_frame_id': insert.child_frame_id,
+            'transform': trans,
+        }
+        transforms.append(tf_msg)
+
+    return {'transforms': transforms}
 
 
 def tf_remove(cfg: SettingTF, msg: DecodedMessageTuple):