diff --git a/mad_icp/apps/mad_icp.py b/mad_icp/apps/mad_icp.py index b95a315..5056d30 100755 --- a/mad_icp/apps/mad_icp.py +++ b/mad_icp/apps/mad_icp.py @@ -39,6 +39,7 @@ from datetime import datetime 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.kitti_reader import KittiReader from mad_icp.apps.utils.visualizer import Visualizer from mad_icp.configurations.datasets.dataset_configurations import DatasetConfiguration_lut @@ -53,12 +54,14 @@ class InputDataInterface(str, Enum): kitti = "kitti", ros1 = "ros1" + ros2 = "ros2" # Can insert additional conversion formats InputDataInterface_lut = { InputDataInterface.kitti: KittiReader, - InputDataInterface.ros1: Ros1Reader + InputDataInterface.ros1: Ros1Reader, + InputDataInterface.ros2: Ros2Reader } @@ -78,7 +81,6 @@ def main(data_path: Annotated[ bool, typer.Option(help="if true anytime realtime", show_default=True)] = False, noviz: Annotated[ bool, typer.Option(help="if true visualizer on", show_default=True)] = False) -> None: - if not data_path.exists(): console.print(f"[red] {data_path} does not exist!") sys.exit(-1) @@ -92,9 +94,14 @@ def main(data_path: Annotated[ visualizer = Visualizer() reader_type = InputDataInterface.kitti + print(data_path) + print(list(data_path.glob("*.db3"))) if len(list(data_path.glob("*.bag"))) != 0: - console.print("[yellow] The dataset is in rosbag format") + 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") + reader_type = InputDataInterface.ros2 else: console.print("[yellow] The dataset is in kitti format") @@ -117,7 +124,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 == InputDataInterface.ros1: + if reader_type in [InputDataInterface.ros1, InputDataInterface.ros2]: topic = data_cf["rosbag_topic"] lidar_to_base = np.array(data_cf["lidar_to_base"]) diff --git a/mad_icp/apps/utils/ros2_reader.py b/mad_icp/apps/utils/ros2_reader.py new file mode 100644 index 0000000..0e6d0fa --- /dev/null +++ b/mad_icp/apps/utils/ros2_reader.py @@ -0,0 +1,90 @@ +# 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 Ros2Reader: + 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: + """ + topic = kwargs.pop('topic') + try: + from rosbags.highlevel import AnyReader + except ModuleNotFoundError: + print("Rosbags library not installed, run 'pip install -U rosbags'") + sys.exit(-1) + + + self.bag = AnyReader([data_dir]) + + self.bag.open() + connection = self.bag.connections + + if not topic: + raise Exception("You have to specify a topic") + + print("Reading the following topic: ", topic) + connection = [x for x in self.bag.connections if x.topic == topic] + self.msgs = self.bag.messages(connections=connection) + + self.min_range = min_range + self.max_range = max_range + self.topic = topic + self.num_messages = self.bag.topics[topic].msgcount + + def __len__(self): + return self.num_messages + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if hasattr(self, "bag"): + self.bag.close() + + def __getitem__(self, item) -> Tuple[float, Tuple[np.ndarray, np.ndarray]]: + connection, timestamp, rawdata = next(self.msgs) + msg = self.bag.deserialize(rawdata, connection.msgtype) + cloud_stamp = msg.header.stamp.sec + msg.header.stamp.nanosec * 1e-9 + points, _ = read_point_cloud( + msg, min_range=self.min_range, max_range=self.max_range) + return timestamp, points diff --git a/pyproject.toml b/pyproject.toml index dc537dd..51a2357 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build" [project] name = "mad-icp" -version = "0.0.3" +version = "0.0.4" description = "It Is All About Matching Data -- Robust and Informed LiDAR Odometry" readme = "README.md" authors = [