diff --git a/python/pyproject.toml b/python/pyproject.toml index 1d3bc7c..d82effd 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "vbr-devkit" -version = "0.0.8" +version = "0.0.9" description = "Development kit for VBR SLAM dataset" readme = "README.md" authors = [ diff --git a/python/vbr_devkit/__init__.py b/python/vbr_devkit/__init__.py index 27fdca4..00ec2dc 100644 --- a/python/vbr_devkit/__init__.py +++ b/python/vbr_devkit/__init__.py @@ -1 +1 @@ -__version__ = "0.0.3" +__version__ = "0.0.9" diff --git a/python/vbr_devkit/datasets/kitti.py b/python/vbr_devkit/datasets/kitti.py index c086e06..eed8198 100644 --- a/python/vbr_devkit/datasets/kitti.py +++ b/python/vbr_devkit/datasets/kitti.py @@ -53,7 +53,7 @@ def __init__(self, data_dir: Path, topic: str, format_fn=None): self.metadata["topic"] = topic self.metadata["num_messages"] = 0 - def push_back(self, data: Union[PointCloudXf, Image, Imu], timestamp): + def push_back(self, data: Union[PointCloudXf, Image, Imu], timestamp, *args, **kwargs): if "msg_type" not in self.metadata: self.metadata["msg_type"] = data.__class__.__name__ @@ -61,11 +61,11 @@ def push_back(self, data: Union[PointCloudXf, Image, Imu], timestamp): raise RuntimeError( f"TopicHandler is bound to {self.metadata['msg_type']}. Cannot handle data of type {type(data)}") - self.save_fn[self.metadata["msg_type"]](data, timestamp) + self.save_fn[self.metadata["msg_type"]](data, timestamp, *args, **kwargs) self.timestamps.append(timestamp) self.metadata["num_messages"] += 1 - def _save_cloud(self, data: PointCloudXf, timestamp): + def _save_cloud(self, data: PointCloudXf, timestamp, *args, **kwargs): dest_path = self.data_f / Path(self.format_fn(self.metadata["num_messages"]) + ".bin") # Save fields to metadata to recover it later. # We assume fields to remain constant through data of this topic @@ -78,16 +78,28 @@ def _save_cloud(self, data: PointCloudXf, timestamp): with open(self.data_f / ".dtype.pkl", "wb") as f: pickle.dump(data.points.dtype, f) - data.points.tofile(dest_path) + if "pcloud_kitti_format" in kwargs: + if kwargs.get("pcloud_kitti_format"): + clip_points = np.stack([data.points["x"], data.points["y"], data.points["z"], data.points["intensity"]], + axis=1) + clip_points.tofile(dest_path) + else: + data.points.tofile(dest_path) - def _save_image(self, data: Image, timestamp: float): + def _save_image(self, data: Image, timestamp: float, *args, **kwargs): dest_path = self.data_f / Path(self.format_fn(self.metadata["num_messages"]) + ".png") + + if "rgb_convert" in kwargs: + if kwargs.get("rgb_convert"): + data.image = cv2.cvtColor(data.image, cv2.COLOR_BAYER_RG2RGB) + data.encoding = "rgb8" + if not "encoding" in self.metadata.keys(): self.metadata["encoding"] = data.encoding cv2.imwrite(str(dest_path), data.image) - def _save_imu(self, data: Imu, timestamp: float): + def _save_imu(self, data: Imu, timestamp: float, *args, **kwargs): if not self.imu_dest: self.imu_dest = (self.data_f / "imu.txt").open("w") self.imu_dest.write(IMU_CSV_HEADER) @@ -109,10 +121,12 @@ def close(self): class KittiWriter: - def __init__(self, data_dir: Path): + def __init__(self, data_dir: Path, rgb_convert: bool = True, pcloud_kitti_format: bool = True, *args, **kwargs): data_dir.mkdir(parents=True, exist_ok=True) self.destination_dir = data_dir self.data_handles = {} + self.rgb_convert = rgb_convert + self.pcloud_kitti_format = pcloud_kitti_format def __enter__(self): return self @@ -127,8 +141,9 @@ def publish(self, timestamp, topic: str, message: Union[PointCloudXf, Image, Imu self.data_handles[topic] = KittiTopicHandler(self.destination_dir / Path(handle_dir), topic, lambda x: f"{x:010d}") - self.data_handles[topic].push_back(message, timestamp) + self.data_handles[topic].push_back(message, timestamp, rgb_convert=self.rgb_convert, + pcloud_kitti_format=self.pcloud_kitti_format) def __exit__(self, exc_type, exc_val, exc_tb): for handle in self.data_handles: - self.data_handles[handle].close() \ No newline at end of file + self.data_handles[handle].close() diff --git a/python/vbr_devkit/tools/run.py b/python/vbr_devkit/tools/run.py index 9553a44..de7a7ee 100644 --- a/python/vbr_devkit/tools/run.py +++ b/python/vbr_devkit/tools/run.py @@ -63,7 +63,16 @@ def convert(to: Annotated[OutputDataInterface, typer.Argument(help="Desired data Path, typer.Argument(help="Input bag or directory containing multiple bags", show_default=False)], output_dir: Annotated[ Path, typer.Argument(help="Output directory in which the data will be stored", show_default=False)], - ) -> None: + rgb_conversion: Annotated[ + bool, typer.Option( + help="Enable BayerRG8->RGB conversion during conversion in KITTI format." + " Disable this flag to reduce the memory footprint of the converted folder.", + show_default=True)] = True, + reduce_pcloud: Annotated[ + bool, typer.Option( + help="Preserve only channels during PointCloud conversion in KITTI format. " + "Allows compatibility with KITTI readers but removes extra LiDAR channels", + show_default=True)] = True) -> None: console.print(f"Converting {input_dir} to {to} format at {output_dir}") if to == OutputDataInterface.ros2: if not input_dir.is_dir(): @@ -75,7 +84,8 @@ def convert(to: Annotated[OutputDataInterface, typer.Argument(help="Desired data rosconvert.convert(item, output_dir / item.stem) else: with RosReader(input_dir) as reader: - with OutputDataInterface_lut[to](output_dir) as writer: + with OutputDataInterface_lut[to](output_dir, rgb_convert=rgb_conversion, + pcloud_kitti_format=reduce_pcloud) as writer: for timestamp, topic, message in track(reader, description="Processing..."): writer.publish(timestamp, topic, message) console.print(":tada: Completed")