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 from math import pi from typing import Optional 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 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_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) tmp_file_path.unlink() def create_projection_data( dataset: Dataset, output_path: Path, colormap_name: str, missing_data_color: str, reverse_colormap: bool, horizontal_resolution: int, vertical_scale: int, horizontal_scale: int, roi_angle_start: float, roi_angle_width: float, render_images: bool, ) -> (np.ndarray, Optional[list[Path]]): rendered_images = [] converted_lidar_frames = [] for i, pc in track( enumerate(dataset, 1), description="Creating projections...", total=len(dataset) ): vertical_resolution = int(pc.data["ring"].max() + 1) # Angle calculation implementation # projected_data = pc.data.copy() # projected_data["arctan"] = np.arctan2(projected_data["y"], projected_data["x"]) # projected_data["arctan_normalized"] = 0.5 * (projected_data["arctan"] / pi + 1.0) # projected_data["arctan_scaled"] = projected_data["arctan_normalized"] * horizontal_resolution # #projected_data["horizontal_position"] = np.floor(projected_data["arctan_scaled"]) # projected_data["horizontal_position"] = np.round(projected_data["arctan_scaled"]) # projected_data["normalized_range"] = 1 / np.sqrt( # projected_data["x"] ** 2 + projected_data["y"] ** 2 + projected_data["z"] ** 2 # ) # duplicates = projected_data[projected_data.duplicated(subset=['ring', 'horizontal_position'], keep=False)].sort_values(by=['ring', 'horizontal_position']) # sorted = projected_data.sort_values(by=['ring', 'horizontal_position']) # FIXME: following pivot fails due to duplicates in the data, some points (x, y) are mapped to the same pixel in the projection, have to decide how to handles # these cases # projected_image_data = projected_data.pivot( # index="ring", columns="horizontal_position", values="normalized_range" # ) # projected_image_data = projected_image_data.reindex(columns=range(horizontal_resolution), fill_value=0) # projected_image_data, output_horizontal_resolution = crop_lidar_data_to_roi( # projected_image_data, roi_angle_start, roi_angle_width, horizontal_resolution # ) # create_2d_projection( # projected_image_data, # output_path / f"frame_{i:04d}_projection.png", # output_path / "tmp.png", # colormap_name, # missing_data_color, # reverse_colormap, # horizontal_resolution=output_horizontal_resolution * horizontal_scale, # vertical_resolution=vertical_resolution * vertical_scale, # ) lidar_data = pc.data.copy() lidar_data["horizontal_position"] = ( lidar_data["original_id"] % 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 = lidar_data.reindex( columns=range(horizontal_resolution), fill_value=0 ) lidar_data = lidar_data.reindex(index=range(vertical_resolution), fill_value=0) lidar_data, output_horizontal_resolution = crop_lidar_data_to_roi( lidar_data, roi_angle_start, roi_angle_width, horizontal_resolution ) converted_lidar_frames.append(lidar_data.to_numpy()) if render_images: image_path = create_2d_projection( lidar_data, output_path / f"frame_{i:04d}.png", output_path / "tmp.png", colormap_name, missing_data_color, reverse_colormap, horizontal_resolution=output_horizontal_resolution * horizontal_scale, vertical_resolution=vertical_resolution * vertical_scale, ) rendered_images.append(image_path) projection_data = np.stack(converted_lidar_frames, axis=0) if render_images: return rendered_images, projection_data else: return projection_data 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( "--pointcloud-topic", 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 rendered frames should be written to", ) parser.add_argument( "--output-no-images", action="store_true", help="do not create individual image files for the projection frames", ) parser.add_argument( "--output-no-video", action="store_true", help="do not create a video file from the projection frames", ) parser.add_argument( "--output-no-numpy", action="store_true", help="do not create a numpy file with the projection data", ) parser.add_argument( "--force-generation", action="store_true", help="if used will force the generation even if output already exists", ) 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( "--horizontal-resolution", default=2048, type=positive_int, help="number of horizontal lidar data points", ) 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", ) 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() output_path = args.output_path / args.input_experiment_path.stem output_path.mkdir(parents=True, exist_ok=True) parser.write_config_file( parser.parse_known_args()[0], output_file_paths=[(output_path / "config.yaml").as_posix()], ) # Create temporary folder for images, if outputting images we use the output folder itself as temp folder tmp_path = output_path / "tmp" if args.output_no_images else output_path / "frames" tmp_path.mkdir(parents=True, exist_ok=True) dataset = load_dataset(args.input_experiment_path, args.pointcloud_topic) images = [] if not args.output_no_images or not args.output_no_video: if not args.force_generation and all( (tmp_path / f"frame_{i:04d}.png").exists() for i in range(1, len(dataset) + 1) ): print( f"Skipping image generation for {args.input_experiment_path} as all frames already exist" ) else: projection_data, images = create_projection_data( dataset, tmp_path, args.colormap_name, args.missing_data_color, args.reverse_colormap, args.horizontal_resolution, args.vertical_scale, args.horizontal_scale, args.roi_angle_start, args.roi_angle_width, render_images=True, ) output_numpy_path = (output_path / args.input_experiment_path.stem).with_suffix( ".npy" ) if not args.output_no_numpy: if not args.force_generation and output_numpy_path.exists(): print( f"Skipping numpy file generation for {args.input_experiment_path} as {output_numpy_path} already exists" ) else: if args.output_no_images: projection_data, _ = create_projection_data( dataset, tmp_path, args.colormap_name, args.missing_data_color, args.reverse_colormap, args.horizontal_resolution, args.vertical_scale, args.horizontal_scale, args.roi_angle_start, args.roi_angle_width, render_images=False, ) # processed_range_data.dump(output_numpy_path) np.save(output_numpy_path, projection_data, fix_imports=False) if not args.output_no_video: if ( not args.force_generation and (output_path / args.input_experiment_path.stem) .with_suffix(".mp4") .exists() ): print( f"Skipping video generation for {args.input_experiment_path} as {output_path / args.input_experiment_path.stem}.mp4 already exists" ) else: input_images_pattern = f"{tmp_path}/frame_%04d.png" create_video_from_images( input_images_pattern, (output_path / args.input_experiment_path.stem).with_suffix(".mp4"), calculate_average_frame_rate(dataset), ) if args.output_no_images: for image in images: image.unlink() tmp_path.rmdir() return 0 if __name__ == "__main__": exit(main())