Fix dataloaders (#76)

* fix dataloaders to return [frame, timestamp] always

* remove unnecessary for loop

* fix formatting

* bump kiss-icp required version to allow proper deskewing

Author: saurabh1002

cpp/map_closures/MapClosures.cpp

@@ -84,10 +84,6 @@ void MapClosures::MatchAndAddToDatabase(const int id,
     self_matches.reserve(orb_keypoints.size());
     self_matcher.knnMatch(orb_descriptors, orb_descriptors, self_matches, 2);
 
-    std::for_each(orb_keypoints.begin(), orb_keypoints.end(), [&](cv::KeyPoint &keypoint) {
-        keypoint.pt.x = keypoint.pt.x + static_cast<float>(density_map.lower_bound.y());
-        keypoint.pt.y = keypoint.pt.y + static_cast<float>(density_map.lower_bound.x());
-    });
     density_maps_.emplace(id, std::move(density_map));
     ground_alignments_.emplace(id, T_ground);
 
@@ -96,7 +92,9 @@ void MapClosures::MatchAndAddToDatabase(const int id,
     std::for_each(self_matches.cbegin(), self_matches.cend(), [&](const auto &self_match) {
         if (self_match[1].distance > self_similarity_threshold) {
             const auto index_descriptor = self_match[0].queryIdx;
-            const auto &keypoint = orb_keypoints[index_descriptor];
+            cv::KeyPoint keypoint = orb_keypoints[index_descriptor];
+            keypoint.pt.x = keypoint.pt.x + static_cast<float>(density_map.lower_bound.y());
+            keypoint.pt.y = keypoint.pt.y + static_cast<float>(density_map.lower_bound.x());
             hbst_matchable.emplace_back(
                 new Matchable(keypoint, orb_descriptors.row(index_descriptor), id));
         }

python/map_closures/datasets/apollo.py

@@ -49,7 +49,7 @@ def __len__(self):
         return len(self.scan_files)
 
     def __getitem__(self, idx):
-        return self.get_scan(self.scan_files[idx])
+        return self.get_scan(self.scan_files[idx]), np.array([])
 
     def get_scan(self, scan_file: str):
         points = np.asarray(self.o3d.io.read_point_cloud(scan_file).points, dtype=np.float64)

python/map_closures/datasets/generic.py

@@ -61,8 +61,8 @@ def __getitem__(self, idx):
         return self.read_point_cloud(self.scan_files[idx])
 
     def read_point_cloud(self, file_path: str):
-        points = self._read_point_cloud(file_path)
-        return points.astype(np.float64)
+        points, timestamps = self._read_point_cloud(file_path)
+        return points.astype(np.float64), timestamps.astype(np.float64)
 
     def _get_point_cloud_reader(self):
         """Attempt to guess with try/catch blocks which is the best point cloud reader to use for
@@ -75,34 +75,77 @@ def _get_point_cloud_reader(self):
         # This is easy, the old KITTI format
         if self.file_extension == "bin":
             print("[WARNING] Reading .bin files, the only format supported is the KITTI format")
-            return lambda file: np.fromfile(file, dtype=np.float32).reshape((-1, 4))[:, :3]
+
+            class ReadKITTI:
+                def __call__(self, file):
+                    return np.fromfile(file, dtype=np.float32).reshape((-1, 4))[:, :3], np.array([])
+
+            return ReadKITTI()
 
         first_scan_file = self.scan_files[0]
 
-        # first try trimesh
+        # first try open3d
+        try:
+            import open3d as o3d
+
+            try_pcd = o3d.t.io.read_point_cloud(first_scan_file)
+            if try_pcd.is_empty():
+                # open3d binding does not raise an exception if file is unreadable or extension is not supported
+                raise Exception("Generic Dataloader| Open3d PointCloud file is empty")
+
+            stamps_keys = ["t", "timestamp", "timestamps", "time", "stamps"]
+            stamp_field = None
+            for key in stamps_keys:
+                try:
+                    try_pcd.point[key]
+                    stamp_field = key
+                    print("Generic Dataloader| found timestamps")
+                    break
+                except:
+                    continue
+
+            class ReadOpen3d:
+                def __init__(self, time_field):
+                    self.time_field = time_field
+                    if self.time_field is None:
+                        self.get_timestamps = lambda _: np.array([])
+                    else:
+                        self.get_timestamps = lambda pcd: pcd.point[self.time_field].numpy().ravel()
+
+                def __call__(self, file):
+                    pcd = o3d.t.io.read_point_cloud(file)
+                    points = pcd.point.positions.numpy()
+                    return points, self.get_timestamps(pcd)
+
+            return ReadOpen3d(stamp_field)
+        except:
+            pass
+
         try:
             import trimesh
 
             trimesh.load(first_scan_file)
-            return lambda file: np.asarray(trimesh.load(file).vertices)
+
+            class ReadTriMesh:
+                def __call__(self, file):
+                    return np.asarray(trimesh.load(file).vertices), np.array([])
+
+            return ReadTriMesh()
         except:
             pass
 
-        # then try pyntcloud
         try:
             from pyntcloud import PyntCloud
 
             PyntCloud.from_file(first_scan_file)
-            return lambda file: PyntCloud.from_file(file).points[["x", "y", "z"]].to_numpy()
-        except:
-            pass
 
-        # lastly with open3d
-        try:
-            import open3d as o3d
+            class ReadPynt:
+                def __call__(self, file):
+                    return PyntCloud.from_file(file).points[["x", "y", "z"]].to_numpy(), np.array(
+                        []
+                    )
 
-            o3d.io.read_point_cloud(first_scan_file)
-            return lambda file: np.asarray(o3d.io.read_point_cloud(file).points, dtype=np.float64)
+            return ReadPynt()
         except:
             print("[ERROR], File format not supported")
             sys.exit(1)

python/map_closures/datasets/helipr.py

@@ -48,16 +48,20 @@ def __init__(self, data_dir: Path, sequence: str, *_, **__):
         if self.sequence_id == "Avia":
             self.format_string = "fffBBBL"
             self.intensity_channel = None
+            self.time_channel = 6
         elif self.sequence_id == "Aeva":
             self.format_string = "ffffflBf"
             self.format_string_no_intensity = "ffffflB"
             self.intensity_channel = 7
+            self.time_channel = 5
         elif self.sequence_id == "Ouster":
             self.format_string = "ffffIHHH"
             self.intensity_channel = 3
+            self.time_channel = 4
         elif self.sequence_id == "Velodyne":
             self.format_string = "ffffHf"
             self.intensity_channel = 3
+            self.time_channel = 5
         else:
             print("[ERROR] Unsupported LiDAR Type")
             sys.exit(1)
@@ -66,7 +70,10 @@ def __len__(self):
         return len(self.scan_files)
 
     def __getitem__(self, idx):
-        return self.read_point_cloud(idx)
+        data = self.get_data(idx)
+        points = self.read_point_cloud(data)
+        timestamps = self.read_timestamps(data)
+        return points, timestamps
 
     def get_data(self, idx: int):
         file_path = self.scan_files[idx]
@@ -89,16 +96,17 @@ def get_data(self, idx: int):
         data = np.stack(list_lines)
         return data
 
-    def read_point_cloud(self, idx: int):
-        data = self.get_data(idx)
-        points = data[:, :3]
-        return points.astype(np.float64)
+    def read_timestamps(self, data: np.ndarray) -> np.ndarray:
+        time = data[:, self.time_channel]
+        return (time - time.min()) / (time.max() - time.min())
+
+    def read_point_cloud(self, data: np.ndarray) -> np.ndarray:
+        return data[:, :3]
 
     def load_poses(self, poses_file):
         from pyquaternion import Quaternion
 
         poses = np.loadtxt(poses_file, delimiter=" ")
-        n = poses.shape[0]
 
         xyz = poses[:, 1:4]
         rotations = np.array(

python/map_closures/datasets/kitti.py

@@ -54,7 +54,7 @@ def __len__(self):
         return len(self.scan_files)
 
     def scans(self, idx):
-        return self.read_point_cloud(self.scan_files[idx])
+        return self.read_point_cloud(self.scan_files[idx]), np.array([])
 
     def apply_calibration(self, poses: np.ndarray) -> np.ndarray:
         """Converts from Velodyne to Camera Frame"""

python/map_closures/datasets/mcap.py

@@ -40,14 +40,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=topic)
-        self.read_point_cloud = read_point_cloud
+        self.msgs = self.read_ros2_messages(self.mcap_file, topics=[self.topic])
+        self.timestamps = []
         self.use_global_visualizer = True
 
     def __del__(self):
@@ -56,6 +60,7 @@ def __del__(self):
 
     def __getitem__(self, idx):
         msg = next(self.msgs).ros_msg
+        self.timestamps.append(self.stamp_to_sec(msg.header.stamp))
         return self.read_point_cloud(msg)
 
     def __len__(self):
@@ -68,6 +73,18 @@ 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
+
+    def get_frames_timestamps(self) -> list:
+        return self.timestamps
+
     def check_topic(self, topic: str) -> str:
         # Extract schema id from the .mcap file that encodes the PointCloud2 msg
         schema_id = [

python/map_closures/datasets/mulran.py

@@ -48,7 +48,7 @@ def read_point_cloud(self, file_path: str):
         timestamps = self.get_timestamps()
         if points.shape[0] != timestamps.shape[0]:
             # MuRan has some broken point clouds, just fallback to no timestamps
-            return points.astype(np.float64), np.ones(points.shape[0])
+            return points.astype(np.float64), np.array([])
         return points.astype(np.float64), timestamps
 
     @staticmethod

python/map_closures/datasets/ncd.py

@@ -63,7 +63,7 @@ def getitem(self, scan_file: str):
         timestamps = self.get_timestamps()
         if points.shape[0] != timestamps.shape[0]:
             # MuRan has some broken point clouds, just fallback to no timestamps
-            return points.astype(np.float64), np.ones(points.shape[0])
+            return points.astype(np.float64), np.array([])
         return points.astype(np.float64), timestamps
 
     @staticmethod

python/map_closures/datasets/ouster.py

@@ -22,30 +22,12 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-import glob
 import os
 from typing import Optional
 
 import numpy as np
 
 
-def find_metadata_json(pcap_file: str) -> str:
-    """Attempts to resolve the metadata json file for a provided pcap file."""
-    dir_path, filename = os.path.split(pcap_file)
-    if not filename:
-        return ""
-    if not dir_path:
-        dir_path = os.getcwd()
-    json_candidates = sorted(glob.glob(f"{dir_path}/*.json"))
-    if not json_candidates:
-        return ""
-    prefix_sizes = list(
-        map(lambda p: len(os.path.commonprefix((filename, os.path.basename(p)))), json_candidates)
-    )
-    max_elem = max(range(len(prefix_sizes)), key=lambda i: prefix_sizes[i])
-    return json_candidates[max_elem]
-
-
 class OusterDataloader:
     """Ouster pcap dataloader"""
 
@@ -83,64 +65,42 @@ def __init__(
         """
 
         try:
-            import ouster.pcap as pcap
-            from ouster import client
+            from ouster.sdk import client, open_source
         except ImportError:
-            print(
-                f'[ERROR] ouster-sdk is not installed on your system, run "pip install ouster-sdk"'
-            )
+            print(f'ouster-sdk is not installed on your system, run "pip install ouster-sdk"')
             exit(1)
 
-        # since we import ouster-sdk's client module locally, we keep it locally as well
-        self._client = client
-
         assert os.path.isfile(data_dir), "Ouster pcap dataloader expects an existing PCAP file"
 
         # we expect `data_dir` param to be a path to the .pcap file, so rename for clarity
         pcap_file = data_dir
 
-        metadata_json = meta or find_metadata_json(pcap_file)
-        if not metadata_json:
-            print("[ERROR] Ouster pcap dataloader can't find metadata json file.")
-            exit(1)
-        print("[INFO] Ouster pcap dataloader: using metadata json: ", metadata_json)
+        print("Indexing Ouster pcap to count the scans number ...")
+        source = open_source(str(pcap_file), meta=[meta] if meta else [], index=True)
 
-        self.data_dir = os.path.dirname(data_dir)
+        # since we import ouster-sdk's client module locally, we keep reference
+        # to it locally as well
+        self._client = client
 
-        with open(metadata_json) as json:
-            self._info_json = json.read()
-            self._info = client.SensorInfo(self._info_json)
+        self.data_dir = os.path.dirname(data_dir)
 
         # lookup table for 2D range image projection to a 3D point cloud
-        self._xyz_lut = client.XYZLut(self._info)
+        self._xyz_lut = client.XYZLut(source.metadata)
 
         self._pcap_file = str(data_dir)
 
-        # read pcap file for the first pass to count scans
-        print("[INFO] Pre-reading Ouster pcap to count the scans number ...")
-        self._source = pcap.Pcap(self._pcap_file, self._info)
-        self._scans_num = sum((1 for _ in client.Scans(self._source)))
-        print(f"[INFO] Ouster pcap total scans number:  {self._scans_num}")
+        self._scans_num = len(source)
+        print(f"Ouster pcap total scans number:  {self._scans_num}")
 
         # frame timestamps array
         self._timestamps = np.linspace(0, self._scans_num, self._scans_num, endpoint=False)
 
-        # start Scans iterator for consumption in __getitem__
-        self._source = pcap.Pcap(self._pcap_file, self._info)
-        self._scans_iter = iter(client.Scans(self._source))
-        self._next_idx = 0
+        self._source = source
 
     def __getitem__(self, idx):
-        # we assume that users always reads sequentially and do not
-        # pass idx as for a random access collection
-        assert self._next_idx == idx, (
-            "Ouster pcap dataloader supports only sequential reads. "
-            f"Expected idx: {self._next_idx}, but got {idx}"
-        )
-        scan = next(self._scans_iter)
-        self._next_idx += 1
-
-        self._timestamps[self._next_idx - 1] = 1e-9 * scan.timestamp[0]
+        scan = self._source[idx]
+
+        self._timestamps[idx] = 1e-9 * scan.timestamp[0]
 
         timestamps = np.tile(np.linspace(0, 1.0, scan.w, endpoint=False), (scan.h, 1))