Feature/better bag bindings

ianvs/CMakeLists.txt

@@ -123,14 +123,14 @@ target_link_libraries(transform_file_lookup PUBLIC yaml-cpp::yaml-cpp CLI11::CLI
 # ######################################################################################
 
 pybind11_add_module(
-  _${PROJECT_NAME}_bindings src/bindings/image_encodings.cpp src/bindings/module.cpp
+  _${PROJECT_NAME}_bindings src/bindings/image_encodings.cpp src/bindings/node_init.cpp src/bindings/module.cpp
 )
 target_include_directories(
   _${PROJECT_NAME}_bindings
   PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
          $<INSTALL_INTERFACE:include/${PROJECT_NAME}>
 )
-ament_target_dependencies(_${PROJECT_NAME}_bindings PUBLIC sensor_msgs)
+target_link_libraries(_${PROJECT_NAME}_bindings PUBLIC ${PROJECT_NAME})
 
 include(CheckIPOSupported)
 check_ipo_supported(RESULT IPO_SUPPORTED)

ianvs/ianvs/__init__.py

@@ -1,14 +1,28 @@
-from ianvs._ianvs_bindings import parse_encoding
-import sensor_msgs.msg
-from typing import Optional
+import contextlib
+
+import ianvs._ianvs_bindings as _bindings
+
 import numpy as np
+from sensor_msgs.msg import Image, CompressedImage
 
 
-def get_image(msg: sensor_msgs.msg.Image) -> Optional[np.ndarray]:
+def get_image(msg: Image | CompressedImage) -> np.ndarray | None:
     """Return appropriate image from message if possible."""
-    info = parse_encoding(msg.encoding)
+    if isinstance(msg, CompressedImage):
+        import imageio.v3
+
+        return imageio.v3.imread(msg.data.tobytes())
+
+    info = _bindings.parse_encoding(msg.encoding)
     if info is None:
         return None
 
     shape = (msg.height, msg.width, info.channels)
     return np.squeeze(np.frombuffer(msg.data, dtype=info.dtype_str).reshape(shape))
+
+
+@contextlib.contextmanager
+def init_node_handle(node_name: str):
+    guard = _bindings._NodeInitHolder(node_name)
+    yield guard.name
+    del guard

ianvs/ianvs/bag_reader.py

@@ -0,0 +1,190 @@
+import heapq
+import itertools
+import functools
+from dataclasses import dataclass, field
+from typing import Any
+
+import tqdm
+import tf2_ros
+import rosbag2_py
+from rclpy.time import Time
+from rclpy.serialization import deserialize_message
+from rclpy.logging import get_logger, get_logging_severity_from_string
+from rosidl_runtime_py.utilities import get_message
+
+
+def _all_nonempty(queues):
+    status = [bool(q) for _, q in queues.items()]
+    return functools.reduce(lambda x, y: x and y, status, True)
+
+
+def _get_stamp(msg, bag_stamp_ns):
+    if not hasattr(msg, "header"):
+        return bag_stamp_ns
+
+    return int(msg.header.stamp.sec * 1.0e9) + msg.header.stamp.nanosec
+
+
+def _valid_set(stamps, max_diff_ns):
+    for lhs, rhs in itertools.combinations(stamps.items(), 2):
+        diff_ns = abs(lhs[1] - rhs[1])
+        if diff_ns > max_diff_ns:
+            return False
+
+    return True
+
+
+@dataclass(order=True)
+class _MsgEntry:
+    stamp_ns: int
+    msg: Any = field(compare=False)
+
+
+class ProgressTracker:
+    """Time-based progress bar for iterating through rosbag."""
+
+    def __init__(self, bag, start_time_ns=None, enabled=True):
+        rel_start = start_time_ns if start_time_ns else bag.start_time
+        rel_duration = bag.duration - (rel_start - bag.start_time)
+
+        self.total_s = rel_duration * 1.0e-9
+        self.enabled = enabled
+        self.pbar = None
+        self.last_stamp = None
+
+    def __enter__(self):
+        self.pbar = None
+        self.last_stamp = None
+        if self.enabled:
+            fmt = "{l_bar}{bar}| {n:.2f}/{total:.2f} [s] [{elapsed}<{remaining}, {rate_fmt}{postfix}]"
+            self.pbar = tqdm.tqdm(
+                total=self.total_s, unit=" bag seconds", bar_format=fmt
+            )
+
+        return self
+
+    def update(self, stamp):
+        if self.pbar is not None and self.last_stamp is not None:
+            self.pbar.update((stamp - self.last_stamp) * 1.0e-9)
+
+        self.last_stamp = stamp
+
+    def __exit__(self, exc_type, exc_msg, traceback):
+        if self.pbar is not None:
+            self.pbar.close()
+
+
+class BagReader:
+    """Wrapper around SequentialReader to provide slighter cleaner interface."""
+
+    def __init__(self, bag_path):
+        self._path = bag_path
+        self._bag = rosbag2_py.SequentialReader()
+
+    @property
+    def start_time(self):
+        return self._bag.get_metadata().starting_time.nanoseconds  # type: ignore
+
+    @property
+    def duration(self):
+        return self._bag.get_metadata().duration.nanoseconds  # type: ignore
+
+    @property
+    def topics(self):
+        return [x for x in self._typenames]
+
+    def open(self):
+        level = get_logging_severity_from_string("WARN")
+        logger = get_logger("rosbag2_storage")
+        logger.set_level(level)
+        self._bag.open_uri(str(self._path))
+        self._typenames = {}
+        for info in self._bag.get_all_topics_and_types():  # type: ignore
+            try:
+                self._typenames[info.name] = get_message(info.type)
+            except Exception as e:
+                logger.error(f"Unknown message type: {e}")
+
+    def close(self):
+        self._bag.close()
+
+    def __enter__(self):
+        self.open()
+        return self
+
+    def __exit__(self, *exc):
+        self.close()
+
+    def lookup_static_transform(self, parent, child):
+        """Find the transform between the parent and child."""
+        buffer = tf2_ros.BufferCore()
+        for _, msg, _ in self.read_messages(["/tf_static"], progress=False):
+            for x in msg.transforms:
+                buffer.set_transform_static(x, "rosbag")
+
+        return buffer.lookup_transform_core(parent, child, Time()).transform
+
+    def read_messages(self, topics=None, start_time_ns=None, progress=True):
+        self._bag.seek(0)
+
+        if topics is not None:
+            missing = [x for x in topics if x not in self._typenames]
+            if len(missing) > 0:
+                logger = get_logger("rosbag2_storage")
+                logger.warning(f"Could not find {missing} (available: {self.topics})")
+                return None
+
+            self._bag.set_filter(rosbag2_py.StorageFilter(topics=topics))
+
+        with ProgressTracker(self, start_time_ns, enabled=progress) as tracker:
+            while self._bag.has_next():
+                topic, data, t = self._bag.read_next()
+                if start_time_ns is not None and t < start_time_ns:
+                    continue
+
+                tracker.update(t)
+                yield topic, deserialize_message(data, self._typenames[topic]), t
+
+    def read_synced_messages(
+        self, topics, max_diff_ns=0, start_time_ns=None, progress=True
+    ):
+        logger = get_logger("rosbag2_storage")
+
+        self._bag.seek(0)
+        missing = [x for x in topics if x not in self._typenames]
+        if len(missing) > 0:
+            logger.warning(f"Could not find {missing} (available: {self.topics})")
+            return None
+
+        unknown = set([])
+        queues = {t: list() for t in topics}
+        self._bag.set_filter(rosbag2_py.StorageFilter(topics=topics))
+        with ProgressTracker(self, start_time_ns, enabled=progress) as tracker:
+            while self._bag.has_next():
+                topic, data, t = self._bag.read_next()
+                if start_time_ns is not None and t < start_time_ns:
+                    continue
+
+                tracker.update(t)
+                queue = queues.get(topic)
+                if queue is None and topic not in unknown:
+                    logger.warning(f"Unexpected topic '{topic}'")
+                    unknown.add(topic)
+                    continue
+
+                msg = deserialize_message(data, self._typenames[topic])
+                msg_stamp = _get_stamp(msg, t)
+                heapq.heappush(queue, _MsgEntry(msg_stamp, msg))
+
+                while _all_nonempty(queues):
+                    stamps = {t: q[0].stamp_ns for t, q in queues.items()}
+                    if not _valid_set(stamps, max_diff_ns):
+                        min_topic, _ = min(stamps.items(), key=lambda x: x[1])
+                        heapq.heappop(queues[min_topic])
+                        continue
+
+                    to_return = {t: q[0].msg for t, q in queues.items()}
+                    for t, q in queues.items():
+                        heapq.heappop(q)
+
+                    yield to_return

ianvs/include/ianvs/bindings/node_init.h

@@ -0,0 +1,6 @@
+#pragma once
+#include <pybind11/pybind11.h>
+
+namespace ianvs::bindings::node_init {
+void add(pybind11::module_& m);
+}

ianvs/package.xml

@@ -20,6 +20,8 @@
   <depend>tf2_ros</depend>
   <depend>pybind11-dev</depend>
   <depend>yaml-cpp</depend>
+  <depend>python3-imageio</depend>
+  <depend>python3-tqdm</depend>
   <test_depend>ament_lint_auto</test_depend>
   <test_depend>ament_lint_common</test_depend>
 

ianvs/src/bindings/module.cpp

@@ -1,5 +1,9 @@
 #include <pybind11/pybind11.h>
 
 #include "ianvs/bindings/image_encodings.h"
+#include "ianvs/bindings/node_init.h"
 
-PYBIND11_MODULE(_ianvs_bindings, m) { ::ianvs::bindings::image_encodings::add(m); }
+PYBIND11_MODULE(_ianvs_bindings, m) {
+  ::ianvs::bindings::image_encodings::add(m);
+  ::ianvs::bindings::node_init::add(m);
+}

ianvs/src/bindings/node_init.cpp

@@ -0,0 +1,33 @@
+#include "ianvs/bindings/node_init.h"
+
+#include <string>
+
+#include "ianvs/node_init.h"
+
+namespace py = pybind11;
+
+namespace ianvs::bindings::node_init {
+
+[[nodiscard]] NodeGuard start_node(const std::string& name) {
+  int argc = 1;
+  auto temp = name;
+  char* argv[] = {temp.data()};
+  return init_node(argc, argv, name);
+}
+
+struct NodeInitHolder {
+  explicit NodeInitHolder(const std::string& name);
+
+  NodeGuard guard;
+  const std::string name;
+};
+
+NodeInitHolder::NodeInitHolder(const std::string& _name) : guard(start_node(_name)), name(_name) {}
+
+void add(py::module_& m) {
+  py::class_<NodeInitHolder>(m, "_NodeInitHolder")
+      .def(py::init<std::string>())
+      .def_readonly("name", &NodeInitHolder::name);
+}
+
+}  // namespace ianvs::bindings::node_init