From 44599724edd338fdc6deabbf974cdb927be7cd6e Mon Sep 17 00:00:00 2001 From: Leonardo Brizi Date: Fri, 13 Sep 2024 11:15:42 -0400 Subject: [PATCH] add support to mcap format, fix viewer issue --- mad_icp/apps/mad_icp.py | 16 ++-- mad_icp/apps/utils/mcap_reader.py | 126 ++++++++++++++++++++++++++++++ mad_icp/apps/utils/visualizer.py | 2 +- pyproject.toml | 2 +- 4 files changed, 139 insertions(+), 7 deletions(-) create mode 100644 mad_icp/apps/utils/mcap_reader.py diff --git a/mad_icp/apps/mad_icp.py b/mad_icp/apps/mad_icp.py index 75bfae2..6916dc9 100755 --- a/mad_icp/apps/mad_icp.py +++ b/mad_icp/apps/mad_icp.py @@ -40,6 +40,7 @@ from mad_icp.apps.utils.utils import write_transformed_pose from mad_icp.apps.utils.ros_reader import Ros1Reader from mad_icp.apps.utils.ros2_reader import Ros2Reader +from mad_icp.apps.utils.mcap_reader import McapReader from mad_icp.apps.utils.kitti_reader import KittiReader from mad_icp.apps.utils.visualizer import Visualizer from mad_icp.configurations.datasets.dataset_configurations import DatasetConfiguration_lut @@ -53,15 +54,17 @@ class InputDataInterface(str, Enum): kitti = "kitti", - ros1 = "ros1" - ros2 = "ros2" + ros1 = "ros1", + ros2 = "ros2", + mcap = "mcap" # Can insert additional conversion formats InputDataInterface_lut = { InputDataInterface.kitti: KittiReader, InputDataInterface.ros1: Ros1Reader, - InputDataInterface.ros2: Ros2Reader + InputDataInterface.ros2: Ros2Reader, + InputDataInterface.mcap: McapReader } @@ -99,8 +102,11 @@ def main(data_path: Annotated[ console.print("[yellow] The dataset is in ros bag format") reader_type = InputDataInterface.ros1 elif len(list(data_path.glob("*.db3"))) != 0: - console.print("[yellow] The dataset is in ros2 bag format") + console.print("[yellow] The dataset is in ros2 db3 format") reader_type = InputDataInterface.ros2 + elif os.path.isfile(data_path) and data_path.suffix == ".mcap": + console.print("[yellow] The dataset is in ros2 mcap format") + reader_type = InputDataInterface.mcap else: console.print("[yellow] The dataset is in kitti format") @@ -123,7 +129,7 @@ def main(data_path: Annotated[ # apply_correction = data_cf["apply_correction"] apply_correction = data_cf.get("apply_correction", False) topic = None - if reader_type in [InputDataInterface.ros1, InputDataInterface.ros2]: + if reader_type in [InputDataInterface.ros1, InputDataInterface.ros2, InputDataInterface.mcap]: topic = data_cf["rosbag_topic"] lidar_to_base = np.array(data_cf["lidar_to_base"]) diff --git a/mad_icp/apps/utils/mcap_reader.py b/mad_icp/apps/utils/mcap_reader.py new file mode 100644 index 0000000..7d3b64d --- /dev/null +++ b/mad_icp/apps/utils/mcap_reader.py @@ -0,0 +1,126 @@ +# Copyright 2024 R(obots) V(ision) and P(erception) group +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. + +import os +import sys +from pathlib import Path +from typing import Tuple +import natsort +from mad_icp.apps.utils.point_cloud2 import read_point_cloud +import numpy as np + + +class McapReader: + def __init__(self, data_dir: Path, min_range=0, + max_range=200, *args, **kwargs): + """ + :param data_dir: Directory containing rosbags or path to a rosbag file + :param topics: Topic to read + :param min_range: minimum range for the points + :param max_range: maximum range for the points + :param args: + :param kwargs: + """ + self.topic = kwargs.pop('topic') + + if not self.topic: + raise Exception("You have to specify a topic") + + try: + from mcap.reader import make_reader + from mcap_ros2.reader import read_ros2_messages + except ModuleNotFoundError: + print("mcap package not installed: run 'pip install -U mcap-ros2-support'") + sys.exit(-1) + + print("Reading the following topic: ", self.topic) + + self.min_range = min_range + self.max_range = max_range + + assert os.path.isfile(data_dir), "mcap dataloader expects an existing MCAP file" + mcap_file = str(data_dir) + + self.bag = make_reader(open(mcap_file, "rb")) + self.summary = self.bag.get_summary() + self.topic = self.check_topic(self.topic) + self.num_messages = sum( + count + for (id, count) in self.summary.statistics.channel_message_counts.items() + if self.summary.channels[id].topic == self.topic + ) + self.msgs = read_ros2_messages(mcap_file, topics=[self.topic]) + self.read_point_cloud = read_point_cloud + + def __len__(self): + return self.num_messages + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + return + + def __getitem__(self, item) -> Tuple[float, Tuple[np.ndarray, np.ndarray]]: + msg = next(self.msgs).ros_msg + points, _ = read_point_cloud( + msg, min_range=self.min_range, max_range=self.max_range) + cloud_stamp = msg.header.stamp.sec + msg.header.stamp.nanosec * 1e-9 + return cloud_stamp, points + + def __del__(self): + if hasattr(self, "bag"): + del self.bag + + def check_topic(self, topic: str) -> str: + # Extract schema id from the .mcap file that encodes the PointCloud2 msg + schema_id = [ + schema.id + for schema in self.summary.schemas.values() + if schema.name == "sensor_msgs/msg/PointCloud2" + ][0] + + point_cloud_topics = [ + channel.topic + for channel in self.summary.channels.values() + if channel.schema_id == schema_id + ] + + def print_available_topics_and_exit(): + print(50 * "-") + for t in point_cloud_topics: + print(f"topic : {t}") + print(50 * "-") + sys.exit(1) + + if topic and topic in point_cloud_topics: + return topic + # when user specified the topic check that exists + if topic and topic not in point_cloud_topics: + print(f"Error: Input bag does not contain any topic with this name: {topic}") + print_available_topics_and_exit() \ No newline at end of file diff --git a/mad_icp/apps/utils/visualizer.py b/mad_icp/apps/utils/visualizer.py index 8c35012..910127a 100644 --- a/mad_icp/apps/utils/visualizer.py +++ b/mad_icp/apps/utils/visualizer.py @@ -50,7 +50,7 @@ def __init__(self): self._initialize_visualizer() def _initialize_visualizer(self): - self.vis.create_window() + self.vis.create_window(window_name="MAD-ICP", width=1280, height=720, visible=True) self.vis.add_geometry(self.current) self.vis.add_geometry(self.local_map) self.vis.get_render_option().background_color = BLACK diff --git a/pyproject.toml b/pyproject.toml index 51a2357..eefb5b2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build" [project] name = "mad-icp" -version = "0.0.4" +version = "0.0.5" description = "It Is All About Matching Data -- Robust and Informed LiDAR Odometry" readme = "README.md" authors = [