From 2eb1bf2e05aee347b216a7dccf18cc65d8fd584c Mon Sep 17 00:00:00 2001 From: Jan Kowalczyk Date: Tue, 11 Jun 2024 15:56:33 +0200 Subject: [PATCH] added pickle file for training with lidar range data to output options --- tools/render2d.py | 220 +++++++++++++++++++++++++++++----------------- tools/util.py | 30 +++++-- 2 files changed, 164 insertions(+), 86 deletions(-) diff --git a/tools/render2d.py b/tools/render2d.py index a1b2d8a..370eef6 100644 --- a/tools/render2d.py +++ b/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, - image_pattern_prefix: str, - tmp_files_path: Path, + output_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)}) - 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( + image_data = fill_sparse_data(pc.data, horizontal_resolution).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] + image_data, output_horizontal_resolution = crop_lidar_data_to_roi( + image_data, roi_angle_start, roi_angle_width, horizontal_resolution + ) 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", + output_path / f"frame_{i:04d}.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, - vertical_resolution=(pc.data['ring'].max() + 1) * vertical_scale + horizontal_resolution=output_horizontal_resolution * horizontal_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", - default=Path("./tmp"), + "--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 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", - default=Path("./output/2d_render.mp4"), - type=Path, - help="path rendered video should be written to", + "--output-pickle", + default=True, + type=bool, + help="if the processed data should be saved as a pickle file", ) parser.add_argument( - "--output-images-prefix", - default="2d_render", - type=str, - help="filename prefix for output", + "--skip-existing", + default=True, + type=bool, + 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: - 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) + output_path = args.output_path / args.input_experiment_path.stem + output_path.mkdir(parents=True, exist_ok=True) - if args.output_video: - args.output_video_path.parent.mkdir(parents=True, exist_ok=True) + # Create temporary folder for images, if outputting images we use the output folder itself as temp folder + 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, - args.output_images_prefix, - args.tmp_files_path, + tmp_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: - input_images_pattern = ( - f"{args.tmp_files_path / args.output_images_prefix}_frame_%04d.png" + if args.output_pickle: + output_pickle_path = ( + 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 diff --git a/tools/util.py b/tools/util.py index c513561..48d5f7b 100644 --- a/tools/util.py +++ b/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",