From e4bd3a00b8f30cf5802e0a96a8d132a7046ce249 Mon Sep 17 00:00:00 2001 From: Sina Samangooei Date: Tue, 2 Jun 2020 16:05:56 +0100 Subject: [PATCH] added support for laser.pkl if they exist --- python/pandaset/sensors.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/python/pandaset/sensors.py b/python/pandaset/sensors.py index a3dec3f..664d8db 100644 --- a/python/pandaset/sensors.py +++ b/python/pandaset/sensors.py @@ -132,6 +132,10 @@ def _load_data_file(self, fp: str) -> None: class Lidar(Sensor): + def load(self) -> None: + super().load() + self._load_laser() + @property def _data_file_extension(self) -> str: return 'pkl.gz' @@ -225,6 +229,15 @@ def set_sensor(self, sensor_id: int) -> None: def _load_data_file(self, fp: str) -> DataFrame: return pd.read_pickle(fp) + def _load_laser(self): + lasers_file_location = f'{self._directory}/laser.pkl' + if not os.path.exists(lasers_file_location): + return + lasers = pd.read_pickle(lasers_file_location) + + for laser, df in zip(lasers, self._data): + df.loc[df.index[df['d'] == 0], "laser"] = laser + class Camera(Sensor): @property