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 4a5f86d..35259d3 100644 --- a/ianvs/ianvs/__init__.py +++ b/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 diff --git a/ianvs/ianvs/bag_reader.py b/ianvs/ianvs/bag_reader.py new file mode 100644 index 0000000..6830d0a --- /dev/null +++ b/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 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/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 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