325 lines
9.4 KiB
Python
325 lines
9.4 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
|
|
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 fill_sparse_data(data: DataFrame, horizontal_resolution: int) -> DataFrame:
|
|
complete_original_ids = DataFrame(
|
|
{
|
|
"original_id": np.arange(
|
|
0,
|
|
(data["ring"].max() + 1) * horizontal_resolution,
|
|
dtype=np.uint32,
|
|
)
|
|
}
|
|
)
|
|
data = complete_original_ids.merge(data, on="original_id", how="left")
|
|
data["ring"] = data["original_id"] // horizontal_resolution
|
|
data["horizontal_position"] = data["original_id"] % horizontal_resolution
|
|
return data
|
|
|
|
|
|
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_projection_data(
|
|
dataset: Dataset,
|
|
horizontal_resolution: int,
|
|
roi_angle_start: float,
|
|
roi_angle_width: float,
|
|
) -> list[Path]:
|
|
converted_lidar_frames = []
|
|
|
|
for i, pc in track(
|
|
enumerate(dataset, 1), description="Rendering images...", total=len(dataset)
|
|
):
|
|
lidar_data = fill_sparse_data(pc.data, 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, _ = crop_lidar_data_to_roi(
|
|
lidar_data, roi_angle_start, roi_angle_width, horizontal_resolution
|
|
)
|
|
converted_lidar_frames.append(lidar_data.to_numpy())
|
|
|
|
return np.stack(converted_lidar_frames, axis=0)
|
|
|
|
|
|
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 render_2d_images(
|
|
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,
|
|
) -> list[Path]:
|
|
rendered_images = []
|
|
|
|
for i, pc in track(
|
|
enumerate(dataset, 1), description="Rendering images...", total=len(dataset)
|
|
):
|
|
image_data = fill_sparse_data(pc.data, horizontal_resolution).pivot(
|
|
index="ring", columns="horizontal_position", values="range"
|
|
)
|
|
|
|
image_data, output_horizontal_resolution = crop_lidar_data_to_roi(
|
|
image_data, roi_angle_start, roi_angle_width, horizontal_resolution
|
|
)
|
|
|
|
normalized_data = (image_data - image_data.min().min()) / (
|
|
image_data.max().max() - image_data.min().min()
|
|
)
|
|
image_path = create_2d_projection(
|
|
normalized_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=(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(
|
|
"--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-images",
|
|
type=bool,
|
|
default=True,
|
|
help="if rendered frames should be outputted as images",
|
|
)
|
|
parser.add_argument(
|
|
"--output-video",
|
|
type=bool,
|
|
default=True,
|
|
help="if rendered frames should be outputted as a video",
|
|
)
|
|
parser.add_argument(
|
|
"--output-pickle",
|
|
default=True,
|
|
type=bool,
|
|
help="if the processed data should be saved as a pickle file",
|
|
)
|
|
parser.add_argument(
|
|
"--skip-existing",
|
|
default=True,
|
|
type=bool,
|
|
help="if true will skip rendering existing files",
|
|
)
|
|
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)
|
|
|
|
# Create temporary folder for images, if outputting images we use the output folder itself as temp folder
|
|
tmp_path = output_path / "frames" if args.output_images else output_path / "tmp"
|
|
tmp_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
dataset = load_dataset(args.input_experiment_path, args.pointcloud_topic)
|
|
|
|
images = render_2d_images(
|
|
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,
|
|
)
|
|
|
|
if args.output_pickle:
|
|
output_pickle_path = (
|
|
output_path / args.input_experiment_path.stem
|
|
).with_suffix(".pkl")
|
|
processed_range_data = create_projection_data(
|
|
dataset,
|
|
args.horizontal_resolution,
|
|
args.roi_angle_start,
|
|
args.roi_angle_width,
|
|
)
|
|
processed_range_data.dump(output_pickle_path)
|
|
|
|
if args.output_video:
|
|
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 not args.output_images:
|
|
for image in images:
|
|
image.unlink()
|
|
tmp_path.rmdir()
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit(main())
|