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 import numpy as np matplotlib.use("Agg") import matplotlib.pyplot as plt from util import ( angle, angle_width, positive_int, load_dataset, existing_path, 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, horizontal_resolution: int, vertical_resolution: int, ): fig, ax = plt.subplots(figsize=(float(horizontal_resolution) / 100, float(vertical_resolution) / 100)) 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((horizontal_resolution, vertical_resolution), 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, horizontal_resolution: int, roi_angle_start: float, roi_angle_width: float, vertical_scale: int, horizontal_scale: int, ) -> 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)}) 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: 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: 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", tmp_files_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, vertical_resolution=(pc.data['ring'].max() + 1) * vertical_scale ) 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-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", 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", ) 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, type=positive_int, help="multiplier for vertical scale, for better visualization", ) parser.add_argument( "--horizontal-scale", default=1, type=positive_int, help="multiplier for horizontal scale, for better visualization", ) 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(args.input_experiment_path, args.pointcloud_topic) 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, args.horizontal_resolution, args.roi_angle_start, args.roi_angle_width, args.vertical_scale, args.horizontal_scale, ) 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())