141 lines
4.8 KiB
Python
141 lines
4.8 KiB
Python
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.pyplot as plt
|
|
|
|
from util import (
|
|
load_dataset_from_bag,
|
|
existing_file,
|
|
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,
|
|
):
|
|
fig, ax = plt.subplots(figsize=(20.48, 5.12))
|
|
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((2048, 512), 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,
|
|
) -> 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())
|
|
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,
|
|
)
|
|
|
|
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-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"
|
|
)
|
|
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")
|
|
|
|
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_from_bag(args.input_bag_path)
|
|
|
|
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,
|
|
)
|
|
|
|
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())
|