added pickle file for training with lidar range data to output options

tools/render2d.py

@@ -16,7 +16,9 @@ matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 
 from util import (
-    angle, angle_width, positive_int,
+    angle,
+    angle_width,
+    positive_int,
     load_dataset,
     existing_path,
     create_video_from_images,
@@ -25,6 +27,70 @@ from util import (
 )
 
 
+def fill_sparse_data(data: DataFrame, horizontal_resolution: int) -> DataFrame:
+    complete_original_ids = DataFrame(
+        {
+            "original_id": np.arange(
+                0,
+                (data["ring"].max() + 1) * horizontal_resolution,
+                dtype=np.uint32,
+            )
+        }
+    )
+    data = complete_original_ids.merge(data, on="original_id", how="left")
+    data["ring"] = data["original_id"] // horizontal_resolution
+    data["horizontal_position"] = data["original_id"] % horizontal_resolution
+    return data
+
+
+def crop_lidar_data_to_roi(
+    data: DataFrame,
+    roi_angle_start: float,
+    roi_angle_width: float,
+    horizontal_resolution: int,
+) -> tuple[DataFrame, int]:
+    if roi_angle_width == 360:
+        return data, horizontal_resolution
+
+    roi_index_start = int(horizontal_resolution / 360 * roi_angle_start)
+    roi_index_width = int(horizontal_resolution / 360 * roi_angle_width)
+    roi_index_end = roi_index_start + roi_index_width
+
+    if roi_index_end < horizontal_resolution:
+        cropped_data = data.iloc[:, roi_index_start:roi_index_end]
+    else:
+        roi_index_end = roi_index_end - horizontal_resolution
+        cropped_data = data.iloc[:, roi_index_end:roi_index_start]
+
+    return cropped_data, roi_index_width
+
+
+def create_projection_data(
+    dataset: Dataset,
+    horizontal_resolution: int,
+    roi_angle_start: float,
+    roi_angle_width: float,
+) -> list[Path]:
+    converted_lidar_frames = []
+
+    for i, pc in track(
+        enumerate(dataset, 1), description="Rendering images...", total=len(dataset)
+    ):
+        lidar_data = fill_sparse_data(pc.data, horizontal_resolution)
+        lidar_data["normalized_range"] = 1 / np.sqrt(
+            lidar_data["x"] ** 2 + lidar_data["y"] ** 2 + lidar_data["z"] ** 2
+        )
+        lidar_data = lidar_data.pivot(
+            index="ring", columns="horizontal_position", values="normalized_range"
+        )
+        lidar_data, _ = crop_lidar_data_to_roi(
+            lidar_data, roi_angle_start, roi_angle_width, horizontal_resolution
+        )
+        converted_lidar_frames.append(lidar_data.to_numpy())
+
+    return np.stack(converted_lidar_frames, axis=0)
+
+
 def create_2d_projection(
     df: DataFrame,
     output_file_path: Path,
@@ -35,7 +101,9 @@ def create_2d_projection(
     horizontal_resolution: int,
     vertical_resolution: int,
 ):
-    fig, ax = plt.subplots(figsize=(float(horizontal_resolution) / 100, float(vertical_resolution) / 100))
+    fig, ax = plt.subplots(
+        figsize=(float(horizontal_resolution) / 100, float(vertical_resolution) / 100)
+    )
     ax.imshow(
         df,
         cmap=get_colormap_with_special_missing_color(
@@ -48,60 +116,50 @@ def create_2d_projection(
     plt.savefig(tmp_file_path, dpi=100, bbox_inches="tight", pad_inches=0)
     plt.close()
     img = Image.open(tmp_file_path)
-    img_resized = img.resize((horizontal_resolution, vertical_resolution), Image.LANCZOS)
+    img_resized = img.resize(
+        (horizontal_resolution, vertical_resolution), Image.LANCZOS
+    )
     img_resized.save(output_file_path)
+    tmp_file_path.unlink()
 
 
 def render_2d_images(
     dataset: Dataset,
-    output_images_path: Path,
+    output_path: Path,
-    image_pattern_prefix: str,
-    tmp_files_path: Path,
     colormap_name: str,
     missing_data_color: str,
     reverse_colormap: bool,
     horizontal_resolution: int,
-    roi_angle_start: float,
-    roi_angle_width: float,
     vertical_scale: int,
     horizontal_scale: int,
+    roi_angle_start: float,
+    roi_angle_width: float,
 ) -> list[Path]:
     rendered_images = []
 
     for i, pc in track(
         enumerate(dataset, 1), description="Rendering images...", total=len(dataset)
     ):
-        complete_original_ids = DataFrame({'original_id': np.arange(0, (pc.data['ring'].max() + 1) * horizontal_resolution, dtype=np.uint32)})
+        image_data = fill_sparse_data(pc.data, horizontal_resolution).pivot(
-        pc.data = complete_original_ids.merge(pc.data, on='original_id', how='left')
-        pc.data['ring'] = (pc.data['original_id'] // horizontal_resolution) 
-        pc.data["horizontal_position"] = pc.data["original_id"] % horizontal_resolution
-        image_data = pc.data.pivot(
             index="ring", columns="horizontal_position", values="range"
         )
 
-        if roi_angle_width != 360:
+        image_data, output_horizontal_resolution = crop_lidar_data_to_roi(
-            roi_index_start = int(horizontal_resolution / 360 * roi_angle_start)
+            image_data, roi_angle_start, roi_angle_width, horizontal_resolution
-            roi_index_width = int(horizontal_resolution / 360 * roi_angle_width)
+        )
-            roi_index_end = roi_index_start + roi_index_width
-
-            if roi_index_end < horizontal_resolution:
-                image_data = image_data.iloc[:, roi_index_start:roi_index_end]
-            else:
-                roi_index_end = roi_index_end - horizontal_resolution
-                image_data = image_data.iloc[:, roi_index_end:roi_index_start]
 
         normalized_data = (image_data - image_data.min().min()) / (
             image_data.max().max() - image_data.min().min()
         )
         image_path = create_2d_projection(
             normalized_data,
-            output_images_path / f"{image_pattern_prefix}_frame_{i:04d}.png",
+            output_path / f"frame_{i:04d}.png",
-            tmp_files_path / "tmp.png",
+            output_path / "tmp.png",
             colormap_name,
             missing_data_color,
             reverse_colormap,
-            horizontal_resolution=(roi_index_width if roi_angle_width != 360 else horizontal_resolution) * horizontal_scale,
+            horizontal_resolution=output_horizontal_resolution * horizontal_scale,
-            vertical_resolution=(pc.data['ring'].max() + 1) * vertical_scale
+            vertical_resolution=(pc.data["ring"].max() + 1) * vertical_scale,
         )
 
         rendered_images.append(image_path)
@@ -120,13 +178,22 @@ def main() -> int:
         "--render-config-file", is_config_file=True, help="yaml config file path"
     )
     parser.add_argument(
-        "--input-experiment-path", required=True, type=existing_path, help="path to experiment. (directly to bag file, to parent folder for mcap)"
+        "--input-experiment-path",
+        required=True,
+        type=existing_path,
+        help="path to experiment. (directly to bag file, to parent folder for mcap)",
     )
     parser.add_argument(
-        "--tmp-files-path",
+        "--pointcloud-topic",
-        default=Path("./tmp"),
+        default="/ouster/points",
+        type=str,
+        help="topic in the ros/mcap bag file containing the point cloud data",
+    )
+    parser.add_argument(
+        "--output-path",
+        default=Path("./output"),
         type=Path,
-        help="path temporary files will be written to",
+        help="path rendered frames should be written to",
     )
     parser.add_argument(
         "--output-images",
@@ -134,12 +201,6 @@ def main() -> int:
         default=True,
         help="if rendered frames should be outputted as images",
     )
-    parser.add_argument(
-        "--output-images-path",
-        default=Path("./output"),
-        type=Path,
-        help="path rendered frames should be written to",
-    )
     parser.add_argument(
         "--output-video",
         type=bool,
@@ -147,16 +208,16 @@ def main() -> int:
         help="if rendered frames should be outputted as a video",
     )
     parser.add_argument(
-        "--output-video-path",
+        "--output-pickle",
-        default=Path("./output/2d_render.mp4"),
+        default=True,
-        type=Path,
+        type=bool,
-        help="path rendered video should be written to",
+        help="if the processed data should be saved as a pickle file",
     )
     parser.add_argument(
-        "--output-images-prefix",
+        "--skip-existing",
-        default="2d_render",
+        default=True,
-        type=str,
+        type=bool,
-        help="filename prefix for output",
+        help="if true will skip rendering existing files",
     )
     parser.add_argument(
         "--colormap-name",
@@ -176,30 +237,12 @@ def main() -> int:
         type=bool,
         help="if colormap should be reversed",
     )
-    parser.add_argument(
-        "--pointcloud-topic",
-        default="/ouster/points",
-        type=str,
-        help="topic in the ros/mcap bag file containing the point cloud data",
-    )
     parser.add_argument(
         "--horizontal-resolution",
         default=2048,
         type=positive_int,
         help="number of horizontal lidar data points",
     )
-    parser.add_argument(
-        "--roi-angle-start",
-        default=0,
-        type=angle,
-        help="angle where roi starts",
-    )
-    parser.add_argument(
-        "--roi-angle-width",
-        default=360,
-        type=angle_width,
-        help="width of roi in degrees",
-    )
     parser.add_argument(
         "--vertical-scale",
         default=1,
@@ -212,48 +255,67 @@ def main() -> int:
         type=positive_int,
         help="multiplier for horizontal scale, for better visualization",
     )
+    parser.add_argument(
+        "--roi-angle-start",
+        default=0,
+        type=angle,
+        help="angle where roi starts",
+    )
+    parser.add_argument(
+        "--roi-angle-width",
+        default=360,
+        type=angle_width,
+        help="width of roi in degrees",
+    )
 
     args = parser.parse_args()
 
-    if args.output_images:
+    output_path = args.output_path / args.input_experiment_path.stem
-        args.output_images_path.mkdir(parents=True, exist_ok=True)
+    output_path.mkdir(parents=True, exist_ok=True)
-        args.tmp_files_path = args.output_images_path
-    else:
-        args.tmp_files_path.mkdir(parents=True, exist_ok=True)
 
-    if args.output_video:
+    # Create temporary folder for images, if outputting images we use the output folder itself as temp folder
-        args.output_video_path.parent.mkdir(parents=True, exist_ok=True)
+    tmp_path = output_path / "frames" if args.output_images else output_path / "tmp"
+    tmp_path.mkdir(parents=True, exist_ok=True)
 
     dataset = load_dataset(args.input_experiment_path, args.pointcloud_topic)
 
     images = render_2d_images(
         dataset,
-        args.tmp_files_path,
+        tmp_path,
-        args.output_images_prefix,
-        args.tmp_files_path,
         args.colormap_name,
         args.missing_data_color,
         args.reverse_colormap,
         args.horizontal_resolution,
-        args.roi_angle_start,
-        args.roi_angle_width,
         args.vertical_scale,
         args.horizontal_scale,
+        args.roi_angle_start,
+        args.roi_angle_width,
     )
 
-    if args.output_video:
+    if args.output_pickle:
-        input_images_pattern = (
+        output_pickle_path = (
-            f"{args.tmp_files_path / args.output_images_prefix}_frame_%04d.png"
+            output_path / args.input_experiment_path.stem
+        ).with_suffix(".pkl")
+        processed_range_data = create_projection_data(
+            dataset,
+            args.horizontal_resolution,
+            args.roi_angle_start,
+            args.roi_angle_width,
         )
+        processed_range_data.dump(output_pickle_path)
+
+    if args.output_video:
+        input_images_pattern = f"{tmp_path}/frame_%04d.png"
         create_video_from_images(
             input_images_pattern,
-            args.output_video_path,
+            (output_path / args.input_experiment_path.stem).with_suffix(".mp4"),
             calculate_average_frame_rate(dataset),
         )
 
     if not args.output_images:
         for image in images:
             image.unlink()
+        tmp_path.rmdir()
 
     return 0
 

tools/util.py

@@ -8,14 +8,17 @@ from matplotlib.colors import Colormap
 from matplotlib import colormaps
 
 
-def load_dataset(bag_file_path: Path, pointcloud_topic: str = "/ouster/points") -> Dataset:
+def load_dataset(
+    bag_file_path: Path, pointcloud_topic: str = "/ouster/points"
+) -> Dataset:
     return Dataset.from_file(bag_file_path, topic=pointcloud_topic)
 
 
-
 def calculate_average_frame_rate(dataset: Dataset):
     timestamps = dataset.timestamps
-    time_deltas = [timestamps[i + 1] - timestamps[i] for i in range(len(timestamps) - 1)]
+    time_deltas = [
+        timestamps[i + 1] - timestamps[i] for i in range(len(timestamps) - 1)
+    ]
     average_delta = sum(time_deltas, timedelta()) / len(time_deltas)
     average_frame_rate = 1 / average_delta.total_seconds()
     return average_frame_rate
@@ -38,39 +41,52 @@ def existing_folder(path_string: str) -> Path:
         raise ArgumentTypeError(f"{path} is not a valid folder!")
     return path
 
+
 def existing_path(path_string: str) -> Path:
     path = Path(path_string)
     if not path.exists():
         raise ArgumentTypeError(f"{path} does not exist!")
     return path
 
+
 def positive_int(number_str: str) -> int:
     number_val = int(number_str)
     if number_val < 0:
         raise ArgumentTypeError(f"{number_val} is not a positive integer!")
     return number_val
 
+
 def angle(angle_str: str) -> float:
     angle_val = float(angle_str)
     if angle_val < 0 or angle_val >= 360:
-        raise ArgumentTypeError(f"{angle_val} is not a valid angle! Needs to be in [0, 360)")
+        raise ArgumentTypeError(
+            f"{angle_val} is not a valid angle! Needs to be in [0, 360)"
+        )
     return angle_val
 
+
 def angle_width(angle_str: str) -> float:
     angle_val = float(angle_str)
     if angle_val < 0 or angle_val > 360:
-        raise ArgumentTypeError(f"{angle_val} is not a valid angle width! Needs to be in [0, 360]")
+        raise ArgumentTypeError(
+            f"{angle_val} is not a valid angle width! Needs to be in [0, 360]"
+        )
     return angle_val
 
+
 def get_colormap_with_special_missing_color(
     colormap_name: str, missing_data_color: str = "black", reverse: bool = False
 ) -> Colormap:
-    colormap = colormaps[colormap_name] if not reverse else colormaps[f"{colormap_name}_r"]
+    colormap = (
+        colormaps[colormap_name] if not reverse else colormaps[f"{colormap_name}_r"]
+    )
     colormap.set_bad(missing_data_color)
     return colormap
 
 
-def create_video_from_images(input_images_pattern: str, output_file: Path, frame_rate: int) -> None:
+def create_video_from_images(
+    input_images_pattern: str, output_file: Path, frame_rate: int
+) -> None:
     # Construct the ffmpeg command
     command = [
         "ffmpeg",