File size: 7,873 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 |
from typing import List
import numpy as np
import numpy.typing as npt
import shapely
from nuplan.common.actor_state.ego_state import EgoState
from nuplan.common.actor_state.state_representation import StateSE2, StateVector2D, TimePoint
from nuplan.common.actor_state.vehicle_parameters import VehicleParameters
from navsim.planning.simulation.planner.pdm_planner.utils.pdm_enums import BBCoordsIndex, SE2Index, StateIndex
from navsim.planning.simulation.planner.pdm_planner.utils.pdm_geometry_utils import translate_lon_and_lat
def array_to_state_se2(array: npt.NDArray[np.float64]) -> StateSE2:
"""
Converts array representation to single StateSE2.
:param array: array filled with (x,y,θ)
:return: StateSE2 class
"""
return StateSE2(array[0], array[1], array[2])
# use numpy vectorize function to apply on last dim
array_to_state_se2_vectorize = np.vectorize(array_to_state_se2, signature="(n)->()")
def array_to_states_se2(array: npt.NDArray[np.float64]) -> npt.NDArray[np.object_]:
"""
Converts array representation to StateSE2 over last dim.
:param array: array filled with (x,y,θ) on last dim
:return: array of StateSE2 class
"""
assert array.shape[-1] == len(SE2Index)
return array_to_state_se2_vectorize(array)
def state_se2_to_array(state_se2: StateSE2) -> npt.NDArray[np.float64]:
"""
Converts StateSE2 to array representation.
:param state_se2: class containing (x,y,θ)
:return: array containing (x,y,θ)
"""
array = np.zeros(len(SE2Index), dtype=np.float64)
array[SE2Index.X] = state_se2.x
array[SE2Index.Y] = state_se2.y
array[SE2Index.HEADING] = state_se2.heading
return array
def states_se2_to_array(states_se2: List[StateSE2]) -> npt.NDArray[np.float64]:
"""
Converts list of StateSE2 object to array representation
:param states_se2: list of StateSE2 object's
:return: array representation of states
"""
state_se2_array = np.zeros((len(states_se2), len(SE2Index)), dtype=np.float64)
for i, state_se2 in enumerate(states_se2):
state_se2_array[i] = state_se2_to_array(state_se2)
return state_se2_array
def ego_state_to_state_array(ego_state: EgoState) -> npt.NDArray[np.float64]:
"""
Converts an ego state into an array representation (drops time-stamps and vehicle parameters)
:param ego_state: ego state class
:return: array containing ego state values
"""
state_array = np.zeros(StateIndex.size(), dtype=np.float64)
state_array[StateIndex.STATE_SE2] = ego_state.rear_axle.serialize()
state_array[StateIndex.VELOCITY_2D] = ego_state.dynamic_car_state.rear_axle_velocity_2d.array
state_array[StateIndex.ACCELERATION_2D] = ego_state.dynamic_car_state.rear_axle_acceleration_2d.array
state_array[StateIndex.STEERING_ANGLE] = ego_state.tire_steering_angle
state_array[StateIndex.STEERING_RATE] = ego_state.dynamic_car_state.tire_steering_rate
state_array[StateIndex.ANGULAR_VELOCITY] = ego_state.dynamic_car_state.angular_velocity
state_array[StateIndex.ANGULAR_ACCELERATION] = ego_state.dynamic_car_state.angular_acceleration
return state_array
def ego_states_to_state_array(ego_states: List[EgoState]) -> npt.NDArray[np.float64]:
"""
Converts a list of ego states into an array representation (drops time-stamps and vehicle parameters)
:param ego_state: ego state class
:return: array containing ego state values
"""
state_array = np.array(
[ego_state_to_state_array(ego_state) for ego_state in ego_states],
dtype=np.float64,
)
return state_array
def state_array_to_ego_state(
state_array: npt.NDArray[np.float64],
time_point: TimePoint,
vehicle_parameters: VehicleParameters,
) -> EgoState:
"""
Converts array representation of ego state back to ego state class.
:param state_array: array representation of ego states
:param time_point: time point of state
:param vehicle_parameters: vehicle parameter of ego
:return: nuPlan's EgoState object
"""
return EgoState.build_from_rear_axle(
rear_axle_pose=StateSE2(*state_array[StateIndex.STATE_SE2]),
rear_axle_velocity_2d=StateVector2D(*state_array[StateIndex.VELOCITY_2D]),
rear_axle_acceleration_2d=StateVector2D(*state_array[StateIndex.ACCELERATION_2D]),
tire_steering_angle=state_array[StateIndex.STEERING_ANGLE],
time_point=time_point,
vehicle_parameters=vehicle_parameters,
is_in_auto_mode=True,
angular_vel=state_array[StateIndex.ANGULAR_VELOCITY],
angular_accel=state_array[StateIndex.ANGULAR_ACCELERATION],
tire_steering_rate=state_array[StateIndex.STEERING_RATE],
)
def state_array_to_ego_states(
state_array: npt.NDArray[np.float64],
time_points: List[TimePoint],
vehicle_parameter: VehicleParameters,
) -> List[EgoState]:
"""
Converts array representation of ego states back to list of ego state class.
:param state_array: array representation of ego states
:param time_point: list of time point of state array
:param vehicle_parameters: vehicle parameter of ego
:return: list nuPlan's EgoState object
"""
ego_states_list: List[EgoState] = []
for i, time_point in enumerate(time_points):
state = state_array[i] if i < len(state_array) else state_array[-1]
ego_states_list.append(state_array_to_ego_state(state, time_point, vehicle_parameter))
return ego_states_list
def state_array_to_coords_array(
states: npt.NDArray[np.float64],
vehicle_parameters: VehicleParameters,
) -> npt.NDArray[np.float64]:
"""
Converts multi-dim array representation of ego states to bounding box coordinates
:param state_array: array representation of ego states
:param vehicle_parameters: vehicle parameter of ego
:return: multi-dim array bounding box coordinates
"""
n_batch, n_time, n_states = states.shape
half_length, half_width, rear_axle_to_center = (
vehicle_parameters.half_length,
vehicle_parameters.half_width,
vehicle_parameters.rear_axle_to_center,
)
headings = states[..., StateIndex.HEADING]
cos, sin = np.cos(headings), np.sin(headings)
# calculate ego center from rear axle
rear_axle_to_center_translate = np.stack([rear_axle_to_center * cos, rear_axle_to_center * sin], axis=-1)
ego_centers: npt.NDArray[np.float64] = states[..., StateIndex.POINT] + rear_axle_to_center_translate
coords_array: npt.NDArray[np.float64] = np.zeros((n_batch, n_time, len(BBCoordsIndex), 2), dtype=np.float64)
coords_array[:, :, BBCoordsIndex.CENTER] = ego_centers
coords_array[:, :, BBCoordsIndex.FRONT_LEFT] = translate_lon_and_lat(ego_centers, headings, half_length, half_width)
coords_array[:, :, BBCoordsIndex.FRONT_RIGHT] = translate_lon_and_lat(
ego_centers, headings, half_length, -half_width
)
coords_array[:, :, BBCoordsIndex.REAR_LEFT] = translate_lon_and_lat(ego_centers, headings, -half_length, half_width)
coords_array[:, :, BBCoordsIndex.REAR_RIGHT] = translate_lon_and_lat(
ego_centers, headings, -half_length, -half_width
)
return coords_array
def coords_array_to_polygon_array(
coords: npt.NDArray[np.float64],
) -> npt.NDArray[np.object_]:
"""
Converts multi-dim array of bounding box coords of to polygons
:param coords: bounding box coords (including corners and center)
:return: array of shapely's polygons
"""
# create coords copy and use center point for closed exterior
coords_exterior: npt.NDArray[np.float64] = coords.copy()
coords_exterior[..., BBCoordsIndex.CENTER, :] = coords_exterior[..., BBCoordsIndex.FRONT_LEFT, :]
# load new coordinates into polygon array
polygons = shapely.creation.polygons(coords_exterior)
return polygons
|