129 lines
4.7 KiB
Python
129 lines
4.7 KiB
Python
from configargparse import ArgParser, YAMLConfigFileParser, ArgumentDefaultsRawHelpFormatter
|
|
from sys import exit
|
|
from open3d.visualization.rendering import OffscreenRenderer, MaterialRecord
|
|
from open3d.io import read_pinhole_camera_parameters, write_image
|
|
from open3d.utility import Vector3dVector
|
|
from pathlib import Path
|
|
from pointcloudset import Dataset
|
|
from rich.progress import track
|
|
|
|
import matplotlib.pyplot as plt
|
|
|
|
import numpy as np
|
|
|
|
from util import (
|
|
load_dataset_from_bag,
|
|
existing_file,
|
|
create_video_from_images,
|
|
calculate_average_frame_rate,
|
|
)
|
|
|
|
|
|
def render_3d_images(
|
|
dataset: Dataset,
|
|
camera_config_input_json_path: Path,
|
|
output_images_path: Path,
|
|
image_pattern_prefix,
|
|
) -> list[Path]:
|
|
camera_params = read_pinhole_camera_parameters("saved_camera_settings.json")
|
|
|
|
width, height = 1920, 1080
|
|
renderer = OffscreenRenderer(width, height)
|
|
renderer.setup_camera(
|
|
intrinsic_matrix=camera_params.intrinsic.intrinsic_matrix,
|
|
extrinsic_matrix=camera_params.extrinsic,
|
|
intrinsic_height_px=camera_params.intrinsic.height,
|
|
intrinsic_width_px=camera_params.intrinsic.width,
|
|
)
|
|
renderer.scene.set_background([1, 1, 1, 1])
|
|
|
|
def color_points_by_range(pcd):
|
|
points = np.asarray(pcd.points)
|
|
distances = np.linalg.norm(points, axis=1)
|
|
max_distance = distances.max()
|
|
min_distance = distances.min()
|
|
normalized_distances = (distances - min_distance) / (max_distance - min_distance)
|
|
colors = plt.get_cmap("jet")(normalized_distances)[:, :3]
|
|
pcd.colors = Vector3dVector(colors)
|
|
return pcd
|
|
|
|
rendered_images = []
|
|
|
|
for i, pc in track(enumerate(dataset, 1), description="Rendering images...", total=len(dataset)):
|
|
o3d_pc = pc.to_instance("open3d")
|
|
o3d_pc = color_points_by_range(o3d_pc)
|
|
renderer.scene.add_geometry("point_cloud", o3d_pc, MaterialRecord())
|
|
image_path = output_images_path / f"{image_pattern_prefix}_{i:04d}.png"
|
|
write_image(image_path.as_posix(), renderer.render_to_image())
|
|
renderer.scene.remove_geometry("point_cloud")
|
|
rendered_images.append(image_path)
|
|
|
|
return rendered_images
|
|
|
|
|
|
def main() -> int:
|
|
parser = ArgParser(
|
|
config_file_parser_class=YAMLConfigFileParser,
|
|
default_config_files=["render3d_config.yaml"],
|
|
formatter_class=ArgumentDefaultsRawHelpFormatter,
|
|
description="Render a 3d representation 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(
|
|
"--camera-config-input-json-path",
|
|
default="./saved_camera_settings.json",
|
|
type=existing_file,
|
|
help="path to json file containing camera settings (can be created with the create_camera_settings.py script)",
|
|
)
|
|
|
|
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_3d_images(
|
|
dataset, args.camera_config_input_json_path, args.tmp_files_path, args.output_images_prefix
|
|
)
|
|
|
|
if args.output_video:
|
|
input_images_pattern = f"{args.tmp_files_path / args.output_images_prefix}_%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())
|