From 7486e110eec24e3cc0543c5a279ec518eff0c441 Mon Sep 17 00:00:00 2001 From: Jan Kowalczyk Date: Tue, 11 Jun 2024 11:07:05 +0200 Subject: [PATCH] adopted 2d rendering for vif dataset --- tools/create_camera_settings.py | 4 +- tools/render2d.py | 160 +++++++++++++++++++++++++++----- tools/render3d.py | 4 +- tools/util.py | 29 +++++- 4 files changed, 168 insertions(+), 29 deletions(-) diff --git a/tools/create_camera_settings.py b/tools/create_camera_settings.py index 6ecae88..c0e762b 100644 --- a/tools/create_camera_settings.py +++ b/tools/create_camera_settings.py @@ -7,7 +7,7 @@ from pathlib import Path from util import ( - load_dataset_from_bag, + load_dataset, existing_file, ) @@ -49,7 +49,7 @@ def main() -> int: args = parser.parse_args() print("Creating camera settings!") print("Move the view in the window to the desired camera position" " and then close the window using the ESC key!") - dataset = load_dataset_from_bag(args.input_bag_path) + dataset = load_dataset(args.input_bag_path) open3d_pointcloud = dataset[args.bag_pointcloud_index].to_instance("open3d") create_camera_settings(open3d_pointcloud, args.camera_config_output_json_path) diff --git a/tools/render2d.py b/tools/render2d.py index 3fb1b16..2d3d6db 100644 --- a/tools/render2d.py +++ b/tools/render2d.py @@ -1,15 +1,24 @@ -from configargparse import ArgParser, YAMLConfigFileParser, ArgumentDefaultsRawHelpFormatter +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 ( - load_dataset_from_bag, - existing_file, + angle, angle_width, positive_int, + load_dataset, + existing_path, create_video_from_images, calculate_average_frame_rate, get_colormap_with_special_missing_color, @@ -23,11 +32,15 @@ def create_2d_projection( colormap_name: str, missing_data_color: str, reverse_colormap: bool, + horizontal_resolution: int, + vertical_resolution: int, ): - fig, ax = plt.subplots(figsize=(20.48, 5.12)) + 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), + cmap=get_colormap_with_special_missing_color( + colormap_name, missing_data_color, reverse_colormap + ), aspect="auto", ) ax.axis("off") @@ -35,7 +48,7 @@ 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((2048, 512), Image.LANCZOS) + img_resized = img.resize((horizontal_resolution, vertical_resolution), Image.LANCZOS) img_resized.save(output_file_path) @@ -47,13 +60,38 @@ def render_2d_images( colormap_name: str, missing_data_color: str, reverse_colormap: bool, + horizontal_resolution: int, + roi_angle_start: float, + roi_angle_width: float, + vertical_scale: int, ) -> 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()) + 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", @@ -61,6 +99,8 @@ def render_2d_images( colormap_name, missing_data_color, reverse_colormap, + horizontal_resolution=roi_index_width if roi_angle_width != 360 else horizontal_resolution, + vertical_resolution=(pc.data['ring'].max() + 1) * vertical_scale ) rendered_images.append(image_path) @@ -75,19 +115,35 @@ def main() -> int: 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" + "--render-config-file", is_config_file=True, help="yaml config file path" ) parser.add_argument( - "--output-images", type=bool, default=True, help="if rendered frames should be outputted as images" + "--input-experiment-path", required=True, type=existing_path, help="path to experiment. (directly to bag file, to parent folder for mcap)" ) parser.add_argument( - "--output-images-path", default=Path("./output"), type=Path, help="path rendered frames should be written to" + "--tmp-files-path", + default=Path("./tmp"), + type=Path, + help="path temporary files will be written to", ) parser.add_argument( - "--output-video", type=bool, default=True, help="if rendered frames should be outputted as a video" + "--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", @@ -95,12 +151,60 @@ def main() -> int: 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" + "--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("--reverse-colormap", default=True, type=bool, help="if colormap should be reversed") args = parser.parse_args() @@ -113,7 +217,7 @@ def main() -> int: if args.output_video: args.output_video_path.parent.mkdir(parents=True, exist_ok=True) - dataset = load_dataset_from_bag(args.input_bag_path) + dataset = load_dataset(args.input_experiment_path, args.pointcloud_topic) images = render_2d_images( dataset, @@ -123,11 +227,21 @@ def main() -> int: args.colormap_name, args.missing_data_color, args.reverse_colormap, + args.horizontal_resolution, + args.roi_angle_start, + args.roi_angle_width, + args.vertical_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)) + 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: diff --git a/tools/render3d.py b/tools/render3d.py index 54173b7..10e13d0 100644 --- a/tools/render3d.py +++ b/tools/render3d.py @@ -12,7 +12,7 @@ import matplotlib.pyplot as plt import numpy as np from util import ( - load_dataset_from_bag, + load_dataset, existing_file, create_video_from_images, calculate_average_frame_rate, @@ -107,7 +107,7 @@ def main() -> int: if args.output_video: args.output_video_path.parent.mkdir(parents=True, exist_ok=True) - dataset = load_dataset_from_bag(args.input_bag_path) + dataset = load_dataset(args.input_bag_path) images = render_3d_images( dataset, args.camera_config_input_json_path, args.tmp_files_path, args.output_images_prefix diff --git a/tools/util.py b/tools/util.py index b2bf252..c513561 100644 --- a/tools/util.py +++ b/tools/util.py @@ -1,4 +1,5 @@ from pointcloudset import Dataset +from pointcloudset.io.dataset.ros import dataset_from_ros from pathlib import Path from argparse import ArgumentTypeError from subprocess import run @@ -7,10 +8,11 @@ from matplotlib.colors import Colormap from matplotlib import colormaps -def load_dataset_from_bag(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)] @@ -32,10 +34,33 @@ def existing_folder(path_string: str) -> Path: path = Path(path_string) if not path.exists(): raise ArgumentTypeError(f"{path} does not exist!") - if not path.is_folder(): + if not path.is_dir(): 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)") + 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]") + return angle_val def get_colormap_with_special_missing_color( colormap_name: str, missing_data_color: str = "black", reverse: bool = False