cras_bag_tools: Added filters to update and dump robot model URDF.

Author: peci1

cras_bag_tools/config/add_subtitles.yaml

@@ -4,6 +4,14 @@
 out_format: '{name}.subs{ext}'
 compression: lz4
 filters:
+  - UpdateRobotModel:
+      urdf_file: "marv.urdf"
+      param: "marv/robot_description"
+  - DumpRobotModel:
+      urdf_file: "marv2.urdf"
+      param: "marv/robot_description"
+      remove_comments: True
+      pretty_print: True
   - AddSubtitles:
-      subtitles_file: '/media/data/bags/outdoor-dataset/marv_2024-10-05-12-34-53.small.bag.srt'
-      topic: /events
+      subtitles_file: 'marv_2024-10-05-12-34-53.small.bag.srt'
+      topic: /commentary

cras_bag_tools/package.xml

@@ -26,6 +26,8 @@
   <exec_depend>kdl_parser_py</exec_depend>
   <exec_depend condition="$ROS_PYTHON_VERSION == 2">python-imageio</exec_depend>
   <exec_depend condition="$ROS_PYTHON_VERSION == 3">python3-imageio</exec_depend>
+  <exec_depend condition="$ROS_PYTHON_VERSION == 2">python-lxml</exec_depend>
+  <exec_depend condition="$ROS_PYTHON_VERSION == 3">python3-lxml</exec_depend>
   <exec_depend condition="$ROS_PYTHON_VERSION == 2">python-marisa</exec_depend>
   <exec_depend condition="$ROS_PYTHON_VERSION == 3">python3-marisa</exec_depend>
   <exec_depend condition="$ROS_PYTHON_VERSION == 2">python-matplotlib</exec_depend>

cras_bag_tools/scripts/filter_bag

@@ -8,6 +8,7 @@
 from __future__ import absolute_import, division, print_function
 
 import argparse
+import copy
 
 import os
 import sys
@@ -141,9 +142,20 @@ def filter_bags(bags, out_format, compression, copy_params, filter, default_para
                     except Exception as e:
                         print("Error loading ROS parameters from file %s: %s" % (params_file, str(e)), file=sys.stderr)
 
+                orig_params = copy.deepcopy(params) if params is not None else None
                 with rosbag.Bag(out_bag_path, 'w', compression=compression) as out:
                     bag.read_index()
                     filter_bag(bag, out, filter, params, start_time, end_time, time_ranges)
+
+                if copy_params and params != orig_params:
+                    print('Saving ROS parameters to', out_bag_path + '.params')
+                    try:
+                        with open(out_bag_path + '.params', 'w+') as f:
+                            yaml.safe_dump(params, f)
+                    except Exception as ex:
+                        print('Error saving ROS params to %s: %s' % (out_bag_path + '.params', str(ex)),
+                              file=sys.stderr)
+
         except Exception as e:
             print('Error processing bag file %s: %s' % (bags_path, str(e)), file=sys.stderr)
             import traceback

cras_bag_tools/src/cras_bag_tools/bag_filter.py

@@ -67,6 +67,8 @@ def filter_bag(bags, out, bag_filter=Passthrough(), params=None, start_time=None
     # apply connection filters
     topics = [c.topic for c in bags._get_connections(topics, bag_filter.connection_filter)]  # noqa
 
+    bag_filter.on_filtering_start()
+
     # get stamp from the tuple
     def get_stamp_fn(x):
         return x[2]
@@ -116,4 +118,6 @@ def get_stamp_fn(x):
         _topic, _raw_msg, _stamp, _connection_header = ret[0]
         out.write(_topic, _raw_msg, _stamp, connection_header=_connection_header, raw=True)
 
+    bag_filter.on_filtering_end()
+
     bag_filter.reset()

cras_bag_tools/src/cras_bag_tools/message_filter.py

@@ -5,10 +5,13 @@
 
 from __future__ import absolute_import, division, print_function
 
+import os.path
+import re
 import sys
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import genpy
+import rospkg
 import rospy
 from cras.message_utils import raw_to_msg, msg_to_raw
 from cras.plugin_utils import get_plugin_implementations
@@ -113,6 +116,8 @@ def __init__(self, is_raw, include_topics=None, exclude_topics=None, include_typ
         self._exclude_time_ranges = self._parse_time_ranges(exclude_time_ranges)
         """Time ranges that specify which regions of the bag should be skipped (but passed further)."""
 
+        self.__rospack = None
+
     @staticmethod
     def _parse_time_ranges(ranges):
         if ranges is None:
@@ -161,6 +166,25 @@ def __get_param(self, params, param, default=None):
         key, rest = param.split("/", 1)
         return self.__get_param(params.get(key, None), rest, default)
 
+    def _set_param(self, param, value):
+        """Set parameter `param` to the parameters set by :meth:`set_params`.
+
+        :param str param: Name of the parameter.
+        :param value: The value to set.
+        """
+        self.__set_param(self._params, param, value)
+
+    def __set_param(self, params, param, value):
+        if param.startswith("/"):
+            param = param[1:]
+        if "/" in param:
+            key, rest = param.split("/", 1)
+            if key not in params:
+                params[key] = {}
+            self.__set_param(params[key], rest, value)
+        else:
+            params[param] = value
+
     def __call__(self, *args, **kwargs):
         """Do the filtering.
 
@@ -278,6 +302,39 @@ def extra_time_ranges(self, bags):
         """
         return None
 
+    def __get_rospack(self):
+        if self.__rospack is None:
+            self.__rospack = rospkg.RosPack()
+        return self.__rospack
+
+    def resolve_file(self, filename):
+        """Resolve `filename` relative to the bag set by :meth:`set_bag`.
+
+        :note: This is ideally called from :meth:`on_filtering_start` or :meth:`filter` because earlier, the `_bag`
+               member variable is not set.
+        """
+        matches = re.match(r'\$\(find ([^)]+)\)', filename)
+        if matches is not None:
+            package_path = self.__get_rospack().get_path(matches[1])
+            filename = filename.replace('$(find %s)' % (matches[1],), package_path)
+
+        if self._bag is None or os.path.isabs(filename) or len(self._bag.filename) == 0:
+            return filename
+
+        return os.path.join(os.path.dirname(self._bag.filename), filename)
+
+    def on_filtering_start(self):
+        """This function is called right before the first message is passed to filter().
+
+        :note: Specifically, :meth:`set_params` and :meth:`set_bag` are already called at this stage.
+        :note: :meth:`extra_initial_messages` will be called after calling this method.
+        """
+        pass
+
+    def on_filtering_end(self):
+        """This function is called right after the last message is processed by filter()."""
+        pass
+
     def reset(self):
         """Reset the filter. This should be called e.g. before starting a new bag."""
         self._bag = None
@@ -483,6 +540,15 @@ def filter(self, topic, datatype, data, md5sum, pytype, stamp, header):
         return topic, datatype, data, md5sum, pytype, stamp, header
 
 
+class NoMessageFilter(RawMessageFilter):
+    """
+    Ignore all messages. Good base for message generators or other helpers.
+    """
+
+    def consider_message(self, topic, datatype, stamp, header, is_from_extra_time_range=False):
+        return False
+
+
 class FilterChain(RawMessageFilter):
     """
     A chain of message filters.
@@ -586,6 +652,16 @@ def extra_final_messages(self):
             for m in f.extra_final_messages():
                 yield m
 
+    def on_filtering_start(self):
+        for f in self.filters:
+            f.on_filtering_start()
+        super(FilterChain, self).on_filtering_start()
+
+    def on_filtering_end(self):
+        for f in self.filters:
+            f.on_filtering_end()
+        super(FilterChain, self).on_filtering_end()
+
     def reset(self):
         for f in self.filters:
             f.reset()

cras_bag_tools/src/cras_bag_tools/message_filters.py

@@ -7,14 +7,14 @@
 
 import copy
 import os.path
-from collections import deque
-from enum import Enum
-from typing import Dict, Optional, Union, Tuple, Sequence
-import yaml
-
 import matplotlib.cm as cmap
 import numpy as np
 import sys
+import yaml
+from collections import deque
+from enum import Enum
+from typing import Dict, Optional, Union, Tuple, Sequence
+from lxml import etree
 from numpy.linalg import inv
 from typing import Any, Dict
 
@@ -43,8 +43,8 @@
 from urdf_parser_py import urdf, xml_reflection
 from vision_msgs.msg import Detection2DArray, Detection2D, ObjectHypothesisWithPose
 
-from .message_filter import DeserializedMessageFilter, DeserializedMessageFilterWithTF, Passthrough, RawMessageFilter, \
-    TopicSet
+from .message_filter import DeserializedMessageFilter, DeserializedMessageFilterWithTF, NoMessageFilter, Passthrough, \
+    RawMessageFilter, TopicSet
 
 
 def urdf_error(message):
@@ -1248,7 +1248,7 @@ def _str_params(self):
         return ", ".join(parts)
 
 
-class AddSubtitles(Passthrough):
+class AddSubtitles(NoMessageFilter):
     """Read subtitles from a text file and add them to the output as std_msgs/String messages."""
 
     def __init__(self, subtitles_file, topic, latch=False, time_offset=0.0):
@@ -1265,7 +1265,14 @@ def __init__(self, subtitles_file, topic, latch=False, time_offset=0.0):
         self._latch = latch
         self._time_offset = rospy.Duration(time_offset)
 
+        self._subtitles = None
+
+    def on_filtering_start(self):
+        super(AddSubtitles, self).on_filtering_start()
+        self._start_time = rospy.Time(self._bag.get_start_time()) + self._time_offset
+        subtitles_file = self.resolve_file(self._subtitles_file)
         self._subtitles = self._parse_subtitles(subtitles_file)
+        print("Read", len(self._subtitles), "subtitles from file", subtitles_file)
 
     def _parse_subtitles(self, subtitles_file):
         if subtitles_file.endswith(".srt"):
@@ -1317,10 +1324,6 @@ def parse_time_str(time_str):
 
         return subtitles
 
-    def set_bag(self, bag):
-        super(AddSubtitles, self).set_bag(bag)
-        self._start_time = rospy.Time(bag.get_start_time()) + self._time_offset
-
     def extra_initial_messages(self):
         connection_header = {
             "callerid": "/bag_filter",
@@ -1338,7 +1341,7 @@ def extra_initial_messages(self):
             yield self._topic, msg, abs_time, connection_header
 
     def _str_params(self):
-        params = ["topic=" + self._topic, "subtitles_file=" + self._subtitles_file]
+        params = ["topic=" + self._topic, "subtitles_file=" + self.resolve_file(self._subtitles_file)]
         if self._time_offset != rospy.Duration(0, 0):
             params.append("time_offset=%f" % (self._time_offset.to_sec(),))
         parent_params = super(AddSubtitles, self)._str_params()
@@ -1347,6 +1350,97 @@ def _str_params(self):
         return ",".join(params)
 
 
+class DumpRobotModel(NoMessageFilter):
+    """Read robot model URDF from ROS params and store it in a file."""
+
+    def __init__(self, urdf_file, param="robot_description", remove_comments=False, pretty_print=False,
+                 run_on_start=True):
+        """
+        :param urdf_file: Path to the URDF file. If relative, it will be resolved relative to the bag set by set_bag.
+        :param param: The parameter where robot model should be read.
+        :param remove_comments: Whether comments should be removed from the output URDF file.
+        :param pretty_print: Whether to pretty-print the URDF file (if False, the original layout is preserved).
+        :param run_on_start: If true, the model will be exported before messages are processed. If false, the model will
+                             be exported after processing all messages.
+        """
+        super(DumpRobotModel, self).__init__()
+        self._urdf_file = urdf_file
+        self._param = param
+        self._remove_comments = remove_comments
+        self._pretty_print = pretty_print
+        self._run_on_start = run_on_start
+
+    def on_filtering_start(self):
+        super(DumpRobotModel, self).on_filtering_start()
+
+        if self._run_on_start:
+            self.dump_model()
+
+    def on_filtering_end(self):
+        super(DumpRobotModel, self).on_filtering_end()
+
+        if not self._run_on_start:
+            self.dump_model()
+
+    def dump_model(self):
+        urdf = self._get_param(self._param)
+        if urdf is None:
+            print('Robot model not found on parameter', self._param, file=sys.stderr)
+            return
+
+        if self._remove_comments or self._pretty_print:
+            parser = etree.XMLParser(remove_comments=self._remove_comments, encoding='utf-8')
+            tree = etree.fromstring(urdf.encode('utf-8'), parser=parser)
+            urdf = etree.tostring(tree, encoding='utf-8', xml_declaration=True,
+                                  pretty_print=self._pretty_print).decode('utf-8')
+
+        dest = self.resolve_file(self._urdf_file)
+        with open(dest, 'w+') as f:
+            print(urdf, file=f)
+        print("Robot model from parameter", self._param, "exported to", dest)
+
+    def _str_params(self):
+        params = ["param=" + self._param, "urdf_file=" + self.resolve_file(self._urdf_file)]
+        parent_params = super(DumpRobotModel, self)._str_params()
+        if len(parent_params) > 0:
+            params.append(parent_params)
+        return ",".join(params)
+
+
+class UpdateRobotModel(NoMessageFilter):
+    """Read robot model from URDF file and update it in the ROS parameters."""
+
+    def __init__(self, urdf_file, param="robot_description"):
+        """
+        :param urdf_file: Path to the URDF file. If relative, it will be resolved relative to the bag set by set_bag.
+        :param param: The parameter where robot model should be stored.
+        """
+        super(UpdateRobotModel, self).__init__()
+        self._urdf_file = urdf_file
+        self._param = param
+
+    def on_filtering_start(self):
+        super(UpdateRobotModel, self).on_filtering_start()
+
+        src = self.resolve_file(self._urdf_file)
+        if not os.path.exists(src):
+            print('Cannot find robot URDF file', src, file=sys.stderr)
+            return
+
+        with open(src, 'r', encoding='utf-8') as f:
+            urdf = f.read()
+
+        self._set_param(self._param, urdf)
+        print('Robot model from %s set to parameter %s' % (src, self._param))
+
+    def _str_params(self):
+        params = ["param=" + self._param, "urdf_file=" + self.resolve_file(self._urdf_file)]
+        parent_params = super(UpdateRobotModel, self)._str_params()
+        if len(parent_params) > 0:
+            params.append(parent_params)
+        return ",".join(params)
+
+
 class FixSpotCams(RawMessageFilter):
     """Fix a problem with Spot robot cameras that publish a bit weird message header."""