Gupta/fix serialized dataloaders (#445)

* add reset option to mcap dataloader

* add reset option to rosbag dataloader

* reset timestamps as well

Author: saurabh1002

python/kiss_icp/datasets/mcap.py

@@ -42,15 +42,18 @@ def __init__(self, data_dir: str, topic: str, *_, **__):
         # we expect `data_dir` param to be a path to the .mcap file, so rename for clarity
         assert os.path.isfile(data_dir), "mcap dataloader expects an existing MCAP file"
         self.sequence_id = os.path.basename(data_dir).split(".")[0]
-        mcap_file = str(data_dir)
+        self.mcap_file = str(data_dir)
 
-        self.bag = make_reader(open(mcap_file, "rb"))
+        self.make_reader = make_reader
+        self.read_ros2_messages = read_ros2_messages
+        self.read_point_cloud = read_point_cloud
+
+        self.bag = self.make_reader(open(self.mcap_file, "rb"))
         self.summary = self.bag.get_summary()
         self.topic = self.check_topic(topic)
         self.n_scans = self._get_n_scans()
-        self.msgs = read_ros2_messages(mcap_file, topics=[self.topic])
+        self.msgs = self.read_ros2_messages(self.mcap_file, topics=[self.topic])
         self.timestamps = []
-        self.read_point_cloud = read_point_cloud
         self.use_global_visualizer = True
 
     def __del__(self):
@@ -72,6 +75,11 @@ def _get_n_scans(self) -> int:
             if self.summary.channels[id].topic == self.topic
         )
 
+    def reset(self):
+        self.timestamps = []
+        self.bag = self.make_reader(open(self.mcap_file, "rb"))
+        self.msgs = self.read_ros2_messages(self.mcap_file, topics=[self.topic])
+
     @staticmethod
     def stamp_to_sec(stamp):
         return stamp.sec + float(stamp.nanosec) / 1e9

python/kiss_icp/datasets/rosbag.py

@@ -68,8 +68,8 @@ def __init__(self, data_dir: Path, topic: str, *_, **__):
         self.n_scans = self.bag.topics[self.topic].msgcount
 
         # limit connections to selected topic
-        connections = [x for x in self.bag.connections if x.topic == self.topic]
-        self.msgs = self.bag.messages(connections=connections)
+        self.connections = [x for x in self.bag.connections if x.topic == self.topic]
+        self.msgs = self.bag.messages(connections=self.connections)
         self.timestamps = []
 
         # Visualization Options
@@ -88,6 +88,12 @@ def __getitem__(self, idx):
         msg = self.bag.deserialize(rawdata, connection.msgtype)
         return self.read_point_cloud(msg)
 
+    def reset(self):
+        self.timestamps = []
+        self.bag.close()
+        self.bag.open()
+        self.msgs = self.bag.messages(connections=self.connections)
+
     @staticmethod
     def to_sec(nsec: int):
         return float(nsec) / 1e9