lkllkl's picture
Upload folder using huggingface_hub
da2e2ac verified
raw
history blame
2.58 kB
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]