diff --git a/python/kiss_icp/datasets/ouster.py b/python/kiss_icp/datasets/ouster.py index 4ec3dd11..de242215 100644 --- a/python/kiss_icp/datasets/ouster.py +++ b/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)) diff --git a/python/pyproject.toml b/python/pyproject.toml index 79fbf519..d90826eb 100644 --- a/python/pyproject.toml +++ b/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",