1+ # Copyright 2024 R(obots) V(ision) and P(erception) group
2+ #
3+ # Redistribution and use in source and binary forms, with or without
4+ # modification, are permitted provided that the following conditions are met:
5+ #
6+ # 1. Redistributions of source code must retain the above copyright notice,
7+ # this list of conditions and the following disclaimer.
8+ #
9+ # 2. Redistributions in binary form must reproduce the above copyright notice,
10+ # this list of conditions and the following disclaimer in the documentation
11+ # and/or other materials provided with the distribution.
12+ #
13+ # 3. Neither the name of the copyright holder nor the names of its contributors
14+ # may be used to endorse or promote products derived from this software
15+ # without specific prior written permission.
16+ #
17+ # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
18+ # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
19+ # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
20+ # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
21+ # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
22+ # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
23+ # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
24+ # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
25+ # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
26+ # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
27+ # POSSIBILITY OF SUCH DAMAGE.
28+
29+ import os
30+ import sys
31+ from pathlib import Path
32+ from typing import Tuple
33+ import natsort
34+ from mad_icp .apps .utils .point_cloud2 import read_point_cloud
35+ import numpy as np
36+
37+
38+ class McapReader :
39+ def __init__ (self , data_dir : Path , min_range = 0 ,
40+ max_range = 200 , * args , ** kwargs ):
41+ """
42+ :param data_dir: Directory containing rosbags or path to a rosbag file
43+ :param topics: Topic to read
44+ :param min_range: minimum range for the points
45+ :param max_range: maximum range for the points
46+ :param args:
47+ :param kwargs:
48+ """
49+ self .topic = kwargs .pop ('topic' )
50+
51+ if not self .topic :
52+ raise Exception ("You have to specify a topic" )
53+
54+ try :
55+ from mcap .reader import make_reader
56+ from mcap_ros2 .reader import read_ros2_messages
57+ except ModuleNotFoundError :
58+ print ("mcap package not installed: run 'pip install -U mcap-ros2-support'" )
59+ sys .exit (- 1 )
60+
61+ print ("Reading the following topic: " , self .topic )
62+
63+ self .min_range = min_range
64+ self .max_range = max_range
65+
66+ assert os .path .isfile (data_dir ), "mcap dataloader expects an existing MCAP file"
67+ mcap_file = str (data_dir )
68+
69+ self .bag = make_reader (open (mcap_file , "rb" ))
70+ self .summary = self .bag .get_summary ()
71+ self .topic = self .check_topic (self .topic )
72+ self .num_messages = sum (
73+ count
74+ for (id , count ) in self .summary .statistics .channel_message_counts .items ()
75+ if self .summary .channels [id ].topic == self .topic
76+ )
77+ self .msgs = read_ros2_messages (mcap_file , topics = [self .topic ])
78+ self .read_point_cloud = read_point_cloud
79+
80+ def __len__ (self ):
81+ return self .num_messages
82+
83+ def __enter__ (self ):
84+ return self
85+
86+ def __exit__ (self , exc_type , exc_val , exc_tb ):
87+ return
88+
89+ def __getitem__ (self , item ) -> Tuple [float , Tuple [np .ndarray , np .ndarray ]]:
90+ msg = next (self .msgs ).ros_msg
91+ points , _ = read_point_cloud (
92+ msg , min_range = self .min_range , max_range = self .max_range )
93+ cloud_stamp = msg .header .stamp .sec + msg .header .stamp .nanosec * 1e-9
94+ return cloud_stamp , points
95+
96+ def __del__ (self ):
97+ if hasattr (self , "bag" ):
98+ del self .bag
99+
100+ def check_topic (self , topic : str ) -> str :
101+ # Extract schema id from the .mcap file that encodes the PointCloud2 msg
102+ schema_id = [
103+ schema .id
104+ for schema in self .summary .schemas .values ()
105+ if schema .name == "sensor_msgs/msg/PointCloud2"
106+ ][0 ]
107+
108+ point_cloud_topics = [
109+ channel .topic
110+ for channel in self .summary .channels .values ()
111+ if channel .schema_id == schema_id
112+ ]
113+
114+ def print_available_topics_and_exit ():
115+ print (50 * "-" )
116+ for t in point_cloud_topics :
117+ print (f"topic : { t } )
118+ print (50 * "-" )
119+ sys .exit (1 )
120+
121+ if topic and topic in point_cloud_topics :
122+ return topic
123+ # when user specified the topic check that exists
124+ if topic and topic not in point_cloud_topics :
125+ print (f"Error: Input bag does not contain any topic with this name: { topic } )
126+ print_available_topics_and_exit ()
0 commit comments