diff --git a/src/lidar_visualizer/datasets/ouster.py b/src/lidar_visualizer/datasets/ouster.py index 3c00956..5c752f7 100644 --- a/src/lidar_visualizer/datasets/ouster.py +++ b/src/lidar_visualizer/datasets/ouster.py @@ -23,11 +23,10 @@ # SOFTWARE. import glob +import importlib 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.""" @@ -85,13 +84,25 @@ def __init__( try: import ouster.pcap as pcap from ouster import client + from ouster.client import _utils + from ouster.sdk.examples.colormaps import colorize except ModuleNotFoundError as e: raise ModuleNotFoundError( f'ouster-sdk is not installed on your system, run "pip install ouster-sdk"' ) from e + try: + self.o3d = importlib.import_module("open3d") + except ModuleNotFoundError as e: + raise ModuleNotFoundError( + "Open3D is not installed on your system, to fix this either " + 'run "pip install open3d" ' + "or check https://www.open3d.org/docs/release/getting_started.html" + ) from e # since we import ouster-sdk's client module locally, we keep it locally as well self._client = client + self._colorize = colorize + self._utils = _utils assert os.path.isfile(data_dir), "Ouster pcap dataloader expects an existing PCAP file" @@ -121,14 +132,30 @@ def __init__( self._scans_num = sum((1 for _ in client.Scans(self._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 + def get_color_image(self, scan): + """This function was taken from the Ouster SDK. All rights reserved to Ouster, Inc + + https://github.com/ouster-lidar/ouster_example/blob/master/python/src/ouster/sdk/examples/open3d.py + """ + fields = list(scan.fields) + aes = {} + for field_ind, field in enumerate(fields): + if field in (self._client.ChanField.SIGNAL, self._client.ChanField.SIGNAL2): + aes[field_ind] = self._utils.AutoExposure(0.02, 0.1, 3) + else: + aes[field_ind] = self._utils.AutoExposure() + field_ind = 2 + + # Obtain reflectivity for colorizing the cloud + key = scan.field(fields[field_ind]).astype(float) + aes[field_ind](key) + return self._colorize(key) + def __getitem__(self, idx): # we assume that users always reads sequentially and do not # pass idx as for a random access collection @@ -139,19 +166,18 @@ def __getitem__(self, idx): scan = next(self._scans_iter) self._next_idx += 1 - self._timestamps[self._next_idx - 1] = 1e-9 * scan.timestamp[0] - - timestamps = np.tile(np.linspace(0, 1.0, scan.w, endpoint=False), (scan.h, 1)) - # filtering our zero returns makes it substantially faster for kiss-icp sel_flag = scan.field(self._client.ChanField.RANGE) != 0 - xyz = self._xyz_lut(scan)[sel_flag] - timestamps = timestamps[sel_flag] - return xyz, timestamps + # Extract XYZ and Intensity channels form scan + xyz = self._xyz_lut(scan)[sel_flag] + ref = self.get_color_image(scan)[sel_flag] - def get_frames_timestamps(self): - return self._timestamps + # Build Open3D object + cloud = self.o3d.geometry.PointCloud() + cloud.points = self.o3d.utility.Vector3dVector(xyz.reshape((-1, 3))) + cloud.colors = self.o3d.utility.Vector3dVector(ref.reshape((-1, 3))) + return cloud def __len__(self): return self._scans_num