From 2e3c8c3378f7092d6563a7d9f19b5520c72c3e53 Mon Sep 17 00:00:00 2001 From: Nathan Hughes Date: Sun, 23 Aug 2026 17:17:06 +0000 Subject: [PATCH 1/4] add initial cut at python bag read interface --- ianvs/ianvs/__init__.py | 12 +++ ianvs/ianvs/bag_reader.py | 170 ++++++++++++++++++++++++++++++++++++++ ianvs/package.xml | 2 + 3 files changed, 184 insertions(+) create mode 100644 ianvs/ianvs/bag_reader.py diff --git a/ianvs/ianvs/__init__.py b/ianvs/ianvs/__init__.py index 4a5f86d..6409b15 100644 --- a/ianvs/ianvs/__init__.py +++ b/ianvs/ianvs/__init__.py @@ -3,6 +3,8 @@ from typing import Optional import numpy as np +import imageio.v3 + def get_image(msg: sensor_msgs.msg.Image) -> Optional[np.ndarray]: """Return appropriate image from message if possible.""" @@ -12,3 +14,13 @@ def get_image(msg: sensor_msgs.msg.Image) -> Optional[np.ndarray]: shape = (msg.height, msg.width, info.channels) return np.squeeze(np.frombuffer(msg.data, dtype=info.dtype_str).reshape(shape)) + + +def _parse_image(msg): + if msg is None: + return None + + if "CompressedImage" in type(msg).__name__: + return imageio.v3.imread(msg.data.tobytes()) + + return get_image(msg) diff --git a/ianvs/ianvs/bag_reader.py b/ianvs/ianvs/bag_reader.py new file mode 100644 index 0000000..14fec5b --- /dev/null +++ b/ianvs/ianvs/bag_reader.py @@ -0,0 +1,170 @@ +import heapq +import itertools +import functools + +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] + return functools.reduce(lambda x, y: x and y, status, True) + + +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 + + +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]) + heapq.heappush(queue, msg) + + while _all_nonempty(queues): + stamps = {t: q[0] for t, q in queues.items()} + if not _valid_set(stamps, max_diff_ns): + min_topic, _ = min(stamps.items(), key=lambda x: x[1]) + queues[min_topic].pop(0) + continue + + yield {t: q.pop(0) for t, q in queues.items()} diff --git a/ianvs/package.xml b/ianvs/package.xml index 14a7296..6eb742a 100644 --- a/ianvs/package.xml +++ b/ianvs/package.xml @@ -20,6 +20,8 @@ tf2_ros pybind11-dev yaml-cpp + python3-imageio + python3-tqdm ament_lint_auto ament_lint_common From f772599e38303e72fe6113da4285df77055ab28e Mon Sep 17 00:00:00 2001 From: Nathan Hughes Date: Mon, 24 Aug 2026 01:27:52 +0000 Subject: [PATCH 2/4] fix timestamp sync --- ianvs/ianvs/bag_reader.py | 32 ++++++++++++++++++++++++++------ 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/ianvs/ianvs/bag_reader.py b/ianvs/ianvs/bag_reader.py index 14fec5b..6830d0a 100644 --- a/ianvs/ianvs/bag_reader.py +++ b/ianvs/ianvs/bag_reader.py @@ -1,6 +1,8 @@ import heapq import itertools import functools +from dataclasses import dataclass, field +from typing import Any import tqdm import tf2_ros @@ -12,19 +14,32 @@ def _all_nonempty(queues): - status = [bool(q) for q in 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: + 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.""" @@ -158,13 +173,18 @@ def read_synced_messages( continue msg = deserialize_message(data, self._typenames[topic]) - heapq.heappush(queue, msg) + msg_stamp = _get_stamp(msg, t) + heapq.heappush(queue, _MsgEntry(msg_stamp, msg)) while _all_nonempty(queues): - stamps = {t: q[0] for t, q in queues.items()} + 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]) - queues[min_topic].pop(0) + heapq.heappop(queues[min_topic]) continue - yield {t: q.pop(0) for t, q in queues.items()} + to_return = {t: q[0].msg for t, q in queues.items()} + for t, q in queues.items(): + heapq.heappop(q) + + yield to_return From 4af58d7e39ba0385bc79b29f94b33793015def99 Mon Sep 17 00:00:00 2001 From: Nathan Hughes Date: Mon, 24 Aug 2026 11:20:31 -0400 Subject: [PATCH 3/4] combine and clean up image parsing --- ianvs/ianvs/__init__.py | 22 +++++++--------------- 1 file changed, 7 insertions(+), 15 deletions(-) diff --git a/ianvs/ianvs/__init__.py b/ianvs/ianvs/__init__.py index 6409b15..6dc4a69 100644 --- a/ianvs/ianvs/__init__.py +++ b/ianvs/ianvs/__init__.py @@ -1,26 +1,18 @@ from ianvs._ianvs_bindings import parse_encoding -import sensor_msgs.msg -from typing import Optional +from sensor_msgs.msg import Image, CompressedImage import numpy as np -import imageio.v3 - -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.""" + if isinstance(msg, CompressedImage): + import imageio.v3 + + return imageio.v3.imread(msg.data.tobytes()) + info = 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)) - - -def _parse_image(msg): - if msg is None: - return None - - if "CompressedImage" in type(msg).__name__: - return imageio.v3.imread(msg.data.tobytes()) - - return get_image(msg) From 236bbc447e5badff4afc926766c1cf9de72e3169 Mon Sep 17 00:00:00 2001 From: Nathan Hughes Date: Mon, 24 Aug 2026 15:09:28 -0400 Subject: [PATCH 4/4] add bindings to init node handle --- ianvs/CMakeLists.txt | 4 +-- ianvs/ianvs/__init__.py | 16 +++++++++--- ianvs/include/ianvs/bindings/node_init.h | 6 +++++ ianvs/src/bindings/module.cpp | 6 ++++- ianvs/src/bindings/node_init.cpp | 33 ++++++++++++++++++++++++ 5 files changed, 59 insertions(+), 6 deletions(-) create mode 100644 ianvs/include/ianvs/bindings/node_init.h create mode 100644 ianvs/src/bindings/node_init.cpp diff --git a/ianvs/CMakeLists.txt b/ianvs/CMakeLists.txt index d16a1f3..5cd6086 100644 --- a/ianvs/CMakeLists.txt +++ b/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 $ $ ) -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) diff --git a/ianvs/ianvs/__init__.py b/ianvs/ianvs/__init__.py index 6dc4a69..35259d3 100644 --- a/ianvs/ianvs/__init__.py +++ b/ianvs/ianvs/__init__.py @@ -1,6 +1,9 @@ -from ianvs._ianvs_bindings import parse_encoding -from sensor_msgs.msg import Image, CompressedImage +import contextlib + +import ianvs._ianvs_bindings as _bindings + import numpy as np +from sensor_msgs.msg import Image, CompressedImage def get_image(msg: Image | CompressedImage) -> np.ndarray | None: @@ -10,9 +13,16 @@ def get_image(msg: Image | CompressedImage) -> np.ndarray | None: return imageio.v3.imread(msg.data.tobytes()) - info = parse_encoding(msg.encoding) + 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 diff --git a/ianvs/include/ianvs/bindings/node_init.h b/ianvs/include/ianvs/bindings/node_init.h new file mode 100644 index 0000000..4dd1e31 --- /dev/null +++ b/ianvs/include/ianvs/bindings/node_init.h @@ -0,0 +1,6 @@ +#pragma once +#include + +namespace ianvs::bindings::node_init { +void add(pybind11::module_& m); +} diff --git a/ianvs/src/bindings/module.cpp b/ianvs/src/bindings/module.cpp index e265e1a..4badec6 100644 --- a/ianvs/src/bindings/module.cpp +++ b/ianvs/src/bindings/module.cpp @@ -1,5 +1,9 @@ #include #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); +} diff --git a/ianvs/src/bindings/node_init.cpp b/ianvs/src/bindings/node_init.cpp new file mode 100644 index 0000000..82a2653 --- /dev/null +++ b/ianvs/src/bindings/node_init.cpp @@ -0,0 +1,33 @@ +#include "ianvs/bindings/node_init.h" + +#include + +#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_(m, "_NodeInitHolder") + .def(py::init()) + .def_readonly("name", &NodeInitHolder::name); +} + +} // namespace ianvs::bindings::node_init