added some visualization tools for lidar pointclouds

tools/render2d.py

@@ -0,0 +1,140 @@
+from configargparse import ArgParser, YAMLConfigFileParser, ArgumentDefaultsRawHelpFormatter
+from sys import exit
+from pathlib import Path
+from pointcloudset import Dataset
+from rich.progress import track
+from pandas import DataFrame
+from PIL import Image
+import matplotlib.pyplot as plt
+
+from util import (
+    load_dataset_from_bag,
+    existing_file,
+    create_video_from_images,
+    calculate_average_frame_rate,
+    get_colormap_with_special_missing_color,
+)
+
+
+def create_2d_projection(
+    df: DataFrame,
+    output_file_path: Path,
+    tmp_file_path: Path,
+    colormap_name: str,
+    missing_data_color: str,
+    reverse_colormap: bool,
+):
+    fig, ax = plt.subplots(figsize=(20.48, 5.12))
+    ax.imshow(
+        df,
+        cmap=get_colormap_with_special_missing_color(colormap_name, missing_data_color, reverse_colormap),
+        aspect="auto",
+    )
+    ax.axis("off")
+    fig.subplots_adjust(left=0, right=1, top=1, bottom=0)
+    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((2048, 512), Image.LANCZOS)
+    img_resized.save(output_file_path)
+
+
+def render_2d_images(
+    dataset: Dataset,
+    output_images_path: Path,
+    image_pattern_prefix: str,
+    tmp_files_path: Path,
+    colormap_name: str,
+    missing_data_color: str,
+    reverse_colormap: bool,
+) -> list[Path]:
+    rendered_images = []
+
+    for i, pc in track(enumerate(dataset, 1), description="Rendering images...", total=len(dataset)):
+        pc.data["horizontal_position"] = pc.data["original_id"] % 2048
+        image_data = pc.data.pivot(index="ring", columns="horizontal_position", values="range")
+        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",
+            tmp_files_path / "tmp.png",
+            colormap_name,
+            missing_data_color,
+            reverse_colormap,
+        )
+
+        rendered_images.append(image_path)
+
+    return rendered_images
+
+
+def main() -> int:
+    parser = ArgParser(
+        config_file_parser_class=YAMLConfigFileParser,
+        default_config_files=["render2d_config.yaml"],
+        formatter_class=ArgumentDefaultsRawHelpFormatter,
+        description="Render a 2d projection of a point cloud",
+    )
+    parser.add_argument("--render-config-file", is_config_file=True, help="yaml config file path")
+    parser.add_argument("--input-bag-path", required=True, type=existing_file, help="path to bag file")
+    parser.add_argument(
+        "--tmp-files-path", default=Path("./tmp"), type=Path, help="path temporary files will be written to"
+    )
+    parser.add_argument(
+        "--output-images", type=bool, 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, default=True, help="if rendered frames should be outputted as a video"
+    )
+    parser.add_argument(
+        "--output-video-path",
+        default=Path("./output/2d_render.mp4"),
+        type=Path,
+        help="path rendered video should be written to",
+    )
+    parser.add_argument("--output-images-prefix", default="2d_render", type=str, help="filename prefix for output")
+    parser.add_argument("--colormap-name", default="viridis", type=str, help="name of matplotlib colormap to be used")
+    parser.add_argument(
+        "--missing-data-color", default="black", type=str, help="name of color to be used for missing data"
+    )
+    parser.add_argument("--reverse-colormap", default=True, type=bool, help="if colormap should be reversed")
+
+    args = parser.parse_args()
+
+    if args.output_images:
+        args.output_images_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:
+        args.output_video_path.parent.mkdir(parents=True, exist_ok=True)
+
+    dataset = load_dataset_from_bag(args.input_bag_path)
+
+    images = render_2d_images(
+        dataset,
+        args.tmp_files_path,
+        args.output_images_prefix,
+        args.tmp_files_path,
+        args.colormap_name,
+        args.missing_data_color,
+        args.reverse_colormap,
+    )
+
+    if args.output_video:
+        input_images_pattern = f"{args.tmp_files_path / args.output_images_prefix}_frame_%04d.png"
+        create_video_from_images(input_images_pattern, args.output_video_path, calculate_average_frame_rate(dataset))
+
+    if not args.output_images:
+        for image in images:
+            image.unlink()
+
+    return 0
+
+
+if __name__ == "__main__":
+    exit(main())