File size: 8,305 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 |
from typing import Any, Callable, List, Tuple
import io
from tqdm import tqdm
from PIL import Image
import matplotlib.pyplot as plt
from navsim.agents.abstract_agent import AbstractAgent
from navsim.common.dataclasses import Scene
from navsim.visualization.config import BEV_PLOT_CONFIG, TRAJECTORY_CONFIG, CAMERAS_PLOT_CONFIG
from navsim.visualization.bev import add_configured_bev_on_ax, add_trajectory_to_bev_ax
from navsim.visualization.camera import add_annotations_to_camera_ax, add_lidar_to_camera_ax, add_camera_ax
def configure_bev_ax(ax: plt.Axes) -> plt.Axes:
"""
Configure the plt ax object for birds-eye-view plots
:param ax: matplotlib ax object
:return: configured ax object
"""
margin_x, margin_y = BEV_PLOT_CONFIG["figure_margin"]
ax.set_aspect("equal")
# NOTE: x forward, y sideways
ax.set_xlim(-margin_y / 2, margin_y / 2)
ax.set_ylim(-margin_x / 2, margin_x / 2)
# NOTE: left is y positive, right is y negative
ax.invert_xaxis()
return ax
def configure_ax(ax: plt.Axes) -> plt.Axes:
"""
Configure the ax object for general plotting
:param ax: matplotlib ax object
:return: ax object without a,y ticks
"""
ax.set_xticks([])
ax.set_yticks([])
return ax
def configure_all_ax(ax: List[List[plt.Axes]]) -> List[List[plt.Axes]]:
"""
Iterates through 2D ax list/array to apply configurations
:param ax: 2D list/array of matplotlib ax object
:return: configure axes
"""
for i in range(len(ax)):
for j in range(len(ax[i])):
configure_ax(ax[i][j])
return ax
def plot_bev_frame(scene: Scene, frame_idx: int) -> Tuple[plt.Figure, plt.Axes]:
"""
General plot for birds-eye-view visualization
:param scene: navsim scene dataclass
:param frame_idx: index of selected frame
:return: figure and ax object of matplotlib
"""
fig, ax = plt.subplots(1, 1, figsize=BEV_PLOT_CONFIG["figure_size"])
add_configured_bev_on_ax(ax, scene.map_api, scene.frames[frame_idx])
configure_bev_ax(ax)
configure_ax(ax)
return fig, ax
def plot_bev_with_agent(scene: Scene, agent: AbstractAgent) -> Tuple[plt.Figure, plt.Axes]:
"""
Plots agent and human trajectory in birds-eye-view visualization
:param scene: navsim scene dataclass
:param agent: navsim agent
:return: figure and ax object of matplotlib
"""
human_trajectory = scene.get_future_trajectory()
agent_trajectory = agent.compute_trajectory(scene.get_agent_input())
frame_idx = scene.scene_metadata.num_history_frames - 1
fig, ax = plt.subplots(1, 1, figsize=BEV_PLOT_CONFIG["figure_size"])
add_configured_bev_on_ax(ax, scene.map_api, scene.frames[frame_idx])
add_trajectory_to_bev_ax(ax, human_trajectory, TRAJECTORY_CONFIG["human"])
add_trajectory_to_bev_ax(ax, agent_trajectory, TRAJECTORY_CONFIG["agent"])
configure_bev_ax(ax)
configure_ax(ax)
return fig, ax
def plot_cameras_frame(scene: Scene, frame_idx: int) -> Tuple[plt.Figure, Any]:
"""
Plots 8x cameras and birds-eye-view visualization in 3x3 grid
:param scene: navsim scene dataclass
:param frame_idx: index of selected frame
:return: figure and ax object of matplotlib
"""
frame = scene.frames[frame_idx]
fig, ax = plt.subplots(3, 3, figsize=CAMERAS_PLOT_CONFIG["figure_size"])
add_camera_ax(ax[0, 0], frame.cameras.cam_l0)
add_camera_ax(ax[0, 1], frame.cameras.cam_f0)
add_camera_ax(ax[0, 2], frame.cameras.cam_r0)
add_camera_ax(ax[1, 0], frame.cameras.cam_l1)
add_configured_bev_on_ax(ax[1, 1], scene.map_api, frame)
add_camera_ax(ax[1, 2], frame.cameras.cam_r1)
add_camera_ax(ax[2, 0], frame.cameras.cam_l2)
add_camera_ax(ax[2, 1], frame.cameras.cam_b0)
add_camera_ax(ax[2, 2], frame.cameras.cam_r2)
configure_all_ax(ax)
configure_bev_ax(ax[1, 1])
fig.tight_layout()
fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)
return fig, ax
def plot_cameras_frame_with_lidar(scene: Scene, frame_idx: int) -> Tuple[plt.Figure, Any]:
"""
Plots 8x cameras (including the lidar pc) and birds-eye-view visualization in 3x3 grid
:param scene: navsim scene dataclass
:param frame_idx: index of selected frame
:return: figure and ax object of matplotlib
"""
frame = scene.frames[frame_idx]
fig, ax = plt.subplots(3, 3, figsize=CAMERAS_PLOT_CONFIG["figure_size"])
add_lidar_to_camera_ax(ax[0, 0], frame.cameras.cam_l0, frame.lidar)
add_lidar_to_camera_ax(ax[0, 1], frame.cameras.cam_f0, frame.lidar)
add_lidar_to_camera_ax(ax[0, 2], frame.cameras.cam_r0, frame.lidar)
add_lidar_to_camera_ax(ax[1, 0], frame.cameras.cam_l1, frame.lidar)
add_configured_bev_on_ax(ax[1, 1], scene.map_api, frame)
add_lidar_to_camera_ax(ax[1, 2], frame.cameras.cam_r1, frame.lidar)
add_lidar_to_camera_ax(ax[2, 0], frame.cameras.cam_l2, frame.lidar)
add_lidar_to_camera_ax(ax[2, 1], frame.cameras.cam_b0, frame.lidar)
add_lidar_to_camera_ax(ax[2, 2], frame.cameras.cam_r2, frame.lidar)
configure_all_ax(ax)
configure_bev_ax(ax[1, 1])
fig.tight_layout()
fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)
return fig, ax
def plot_cameras_frame_with_annotations(scene: Scene, frame_idx: int) -> Tuple[plt.Figure, Any]:
"""
Plots 8x cameras (including the bounding boxes) and birds-eye-view visualization in 3x3 grid
:param scene: navsim scene dataclass
:param frame_idx: index of selected frame
:return: figure and ax object of matplotlib
"""
frame = scene.frames[frame_idx]
fig, ax = plt.subplots(3, 3, figsize=CAMERAS_PLOT_CONFIG["figure_size"])
add_annotations_to_camera_ax(ax[0, 0], frame.cameras.cam_l0, frame.annotations)
add_annotations_to_camera_ax(ax[0, 1], frame.cameras.cam_f0, frame.annotations)
add_annotations_to_camera_ax(ax[0, 2], frame.cameras.cam_r0, frame.annotations)
add_annotations_to_camera_ax(ax[1, 0], frame.cameras.cam_l1, frame.annotations)
add_configured_bev_on_ax(ax[1, 1], scene.map_api, frame)
add_annotations_to_camera_ax(ax[1, 2], frame.cameras.cam_r1, frame.annotations)
add_annotations_to_camera_ax(ax[2, 0], frame.cameras.cam_l2, frame.annotations)
add_annotations_to_camera_ax(ax[2, 1], frame.cameras.cam_b0, frame.annotations)
add_annotations_to_camera_ax(ax[2, 2], frame.cameras.cam_r2, frame.annotations)
configure_all_ax(ax)
configure_bev_ax(ax[1, 1])
fig.tight_layout()
fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)
return fig, ax
def frame_plot_to_pil(
callable_frame_plot: Callable[[Scene, int], Tuple[plt.Figure, Any]],
scene: Scene,
frame_indices: List[int],
) -> List[Image.Image]:
"""
Plots a frame according to plotting function and return a list of PIL images
:param callable_frame_plot: callable to plot a single frame
:param scene: navsim scene dataclass
:param frame_indices: list of indices to save
:return: list of PIL images
"""
images: List[Image.Image] = []
for frame_idx in tqdm(frame_indices, desc="Rendering frames"):
fig, ax = callable_frame_plot(scene, frame_idx)
# Creating PIL image from fig
buf = io.BytesIO()
fig.savefig(buf, format="png")
buf.seek(0)
images.append(Image.open(buf).copy())
# close buffer and figure
buf.close()
plt.close(fig)
return images
def frame_plot_to_gif(
file_name: str,
callable_frame_plot: Callable[[Scene, int], Tuple[plt.Figure, Any]],
scene: Scene,
frame_indices: List[int],
duration: float = 500,
) -> None:
"""
Saves a frame-wise plotting function as GIF (hard G)
:param callable_frame_plot: callable to plot a single frame
:param scene: navsim scene dataclass
:param frame_indices: list of indices
:param file_name: file path for saving to save
:param duration: frame interval in ms, defaults to 500
"""
images = frame_plot_to_pil(callable_frame_plot, scene, frame_indices)
images[0].save(file_name, save_all=True, append_images=images[1:], duration=duration, loop=0)
|