Fix Ouster dataloader for 0.11 - 0.12 sdk updates

python/kiss_icp/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,62 +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'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("Ouster pcap dataloader can't find metadata json file.")
-            exit(1)
-        print("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("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)))
+        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))
 

python/pyproject.toml

@@ -49,7 +49,7 @@ dependencies = [
 [project.optional-dependencies]
 all = [
     "open3d>=0.13",
-    "ouster-sdk>=0.7.1",
+    "ouster-sdk>=0.11",
     "pyntcloud",
     "PyYAML",
     "trimesh",