|
from typing import Any, List
|
|
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
|
|
import matplotlib
|
|
from matplotlib import pyplot as plt
|
|
|
|
from navsim.visualization.config import LIDAR_CONFIG
|
|
from navsim.common.enums import LidarIndex
|
|
|
|
|
|
def filter_lidar_pc(lidar_pc: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
|
|
"""
|
|
Filter lidar point cloud according to global configuration
|
|
:param lidar_pc: numpy array of shape (6,n)
|
|
:return: filtered point cloud
|
|
"""
|
|
|
|
pc = lidar_pc.T
|
|
mask = (
|
|
np.ones((len(pc)), dtype=bool)
|
|
& (pc[:, LidarIndex.X] > LIDAR_CONFIG["x_lim"][0])
|
|
& (pc[:, LidarIndex.X] < LIDAR_CONFIG["x_lim"][1])
|
|
& (pc[:, LidarIndex.Y] > LIDAR_CONFIG["y_lim"][0])
|
|
& (pc[:, LidarIndex.Y] < LIDAR_CONFIG["y_lim"][1])
|
|
& (pc[:, LidarIndex.Z] > LIDAR_CONFIG["z_lim"][0])
|
|
& (pc[:, LidarIndex.Z] < LIDAR_CONFIG["z_lim"][1])
|
|
)
|
|
pc = pc[mask]
|
|
return pc.T
|
|
|
|
|
|
def get_lidar_pc_color(
|
|
lidar_pc: npt.NDArray[np.float32], as_hex: bool = False
|
|
) -> List[Any]:
|
|
"""
|
|
Compute color map of lidar point cloud according to global configuration
|
|
:param lidar_pc: numpy array of shape (6,n)
|
|
:param as_hex: whether to return hex values, defaults to False
|
|
:return: list of RGB or hex values
|
|
"""
|
|
|
|
pc = lidar_pc.T
|
|
if LIDAR_CONFIG["color_element"] == "none":
|
|
colors_rgb = np.zeros((len(pc), 3), dtype=np.uin8)
|
|
|
|
else:
|
|
if LIDAR_CONFIG["color_element"] == "distance":
|
|
color_intensities = np.linalg.norm(pc[:, LidarIndex.POSITION], axis=-1)
|
|
else:
|
|
color_element_map = {
|
|
"x": LidarIndex.X,
|
|
"y": LidarIndex.Y,
|
|
"z": LidarIndex.Z,
|
|
"intensity": LidarIndex.INTENSITY,
|
|
"ring": LidarIndex.RING,
|
|
"id": LidarIndex.ID,
|
|
}
|
|
color_intensities = pc[:, color_element_map[LIDAR_CONFIG["color_element"]]]
|
|
|
|
min, max = color_intensities.min(), color_intensities.max()
|
|
norm_intensities = [(value - min) / (max - min) for value in color_intensities]
|
|
colormap = plt.get_cmap("viridis")
|
|
colors_rgb = np.array([colormap(value) for value in norm_intensities])
|
|
colors_rgb = (colors_rgb[:, :3] * 255).astype(np.uint8)
|
|
|
|
assert len(colors_rgb) == len(pc)
|
|
if as_hex:
|
|
return [matplotlib.colors.to_hex(tuple(c / 255.0 for c in rgb)) for rgb in colors_rgb]
|
|
|
|
return [tuple(value) for value in colors_rgb]
|
|
|