mt/tools/render2d.py

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())