Added flags for conversion to preserve pointcloud KITTI format and to automatically de-bayer images

Author: EmanueleGiacomini

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 = [

python/vbr_devkit/__init__.py

@@ -1 +1 @@
-__version__ = "0.0.3"
+__version__ = "0.0.9"

python/vbr_devkit/datasets/kitti.py

@@ -53,19 +53,19 @@ 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__
 
         if self.metadata["msg_type"] != data.__class__.__name__:
             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()
+            self.data_handles[handle].close()

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 <x,y,z,intensity> 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")