adopted 2d rendering for vif dataset

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: