d33c6b1e16
if all datapoints in one ring were missing we would have a shape mismatch for that data so we have to fill the sparse data out in the index axis too
378 lines
12 KiB
Python
378 lines
12 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
|
|
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 = lidar_data.reindex(index=range(vertical_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())
|