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