File size: 2,485 Bytes
7accb91 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
from typing import Any, List
import numpy as np
import numpy.typing as npt
import matplotlib
import matplotlib.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]
|