File size: 7,899 Bytes
663494c |
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 |
import os
import numpy as np
from pyquaternion import Quaternion
from shapely import affinity, ops
from shapely.geometry import LineString, box, MultiPolygon, MultiLineString
from nuplan.common.maps.nuplan_map.map_factory import get_maps_api
from nuplan.common.maps.maps_datatypes import SemanticMapLayer
from nuplan.common.actor_state.oriented_box import OrientedBox
import torch
from nuplan.common.actor_state.state_representation import Point2D, StateSE2
from navsim.planning.scenario_builder.navsim_scenario_utils import tracked_object_types
import mmdet3d_plugin.datasets.utils.calibration as calib_utils
import scipy
import cv2
class DiffusionDriveMap:
def __init__(
self,
config,
map_root,
map_version='nuplan-maps-v1.0',
patch_size=(100, 100), # h, w
map_classes={
'centerline': [SemanticMapLayer.LANE, SemanticMapLayer.LANE_CONNECTOR],
'ped_crossing': [SemanticMapLayer.CROSSWALK],
'road_boundary': [SemanticMapLayer.ROADBLOCK, SemanticMapLayer.INTERSECTION],
# 'sidewalk': [SemanticMapLayer.WALKWAYS]
},
need_merged=['road_boundary'],):
self._config = config
self.map_classes = map_classes
self.patch_size = patch_size
self.need_merged = need_merged
self.MAP_APIS_DICT = {
"us-pa-pittsburgh-hazelwood" : get_maps_api(map_root, map_version, "us-pa-pittsburgh-hazelwood"),
"sg-one-north" : get_maps_api(map_root, map_version, "sg-one-north"),
"us-ma-boston" : get_maps_api(map_root, map_version, "us-ma-boston"),
"us-nv-las-vegas-strip" : get_maps_api(map_root, map_version, "us-nv-las-vegas-strip")
}
def compute_bev_semantic_map(
self, info,
):
"""
Creates sematic map in BEV
:param annotations: annotation dataclass
:param map_api: map interface of nuPlan
:param ego_pose: ego pose in global frame
:return: 2D torch tensor of semantic labels
"""
map_location = info['map_location']
#specific APIS:
map_api = self.MAP_APIS_DICT[map_location]
sdc_loc_global, _ = calib_utils.transform_matrix_to_vector(
info["ego2global"]
)
yaw_global = scipy.spatial.transform.Rotation.from_matrix(
info["ego2global"][:3, :3]
).as_euler("xyz", degrees=False)[-1]
ego_pose = StateSE2(x=sdc_loc_global[0], y=sdc_loc_global[1],heading=yaw_global)
bev_semantic_map = np.zeros(self._config.bev_semantic_frame, dtype=np.int64)
for label, (entity_type, layers) in self._config.bev_semantic_classes.items():
if entity_type == "polygon":
entity_mask = self._compute_map_polygon_mask(map_api, ego_pose, layers)
elif entity_type == "linestring":
entity_mask = self._compute_map_linestring_mask(map_api, ego_pose, layers)
else:
entity_mask = self._compute_box_mask(info, layers)
bev_semantic_map[entity_mask] = label
return torch.Tensor(bev_semantic_map)
def _compute_map_polygon_mask(
self, map_api, ego_pose, layers
) :
"""
Compute binary mask given a map layer class
:param map_api: map interface of nuPlan
:param ego_pose: ego pose in global frame
:param layers: map layers
:return: binary mask as numpy array
"""
map_object_dict = map_api.get_proximal_map_objects(
point=ego_pose.point, radius=self._config.bev_radius, layers=layers
)
map_polygon_mask = np.zeros(self._config.bev_semantic_frame[::-1], dtype=np.uint8)
for layer in layers:
for map_object in map_object_dict[layer]:
polygon = self._geometry_local_coords(map_object.polygon, ego_pose)
exterior = np.array(polygon.exterior.coords).reshape((-1, 1, 2))
exterior = self._coords_to_pixel(exterior)
cv2.fillPoly(map_polygon_mask, [exterior], color=255)
# OpenCV has origin on top-left corner
map_polygon_mask = np.rot90(map_polygon_mask)[::-1]
return map_polygon_mask > 0
def _compute_map_linestring_mask(
self, map_api, ego_pose, layers
):
"""
Compute binary of linestring given a map layer class
:param map_api: map interface of nuPlan
:param ego_pose: ego pose in global frame
:param layers: map layers
:return: binary mask as numpy array
"""
map_object_dict = map_api.get_proximal_map_objects(
point=ego_pose.point, radius=self._config.bev_radius, layers=layers
)
map_linestring_mask = np.zeros(self._config.bev_semantic_frame[::-1], dtype=np.uint8)
for layer in layers:
for map_object in map_object_dict[layer]:
linestring: LineString = self._geometry_local_coords(map_object.baseline_path.linestring, ego_pose)
points = np.array(linestring.coords).reshape((-1, 1, 2))
points = self._coords_to_pixel(points)
cv2.polylines(map_linestring_mask, [points], isClosed=False, color=255, thickness=2)
# OpenCV has origin on top-left corner
map_linestring_mask = np.rot90(map_linestring_mask)[::-1]
return map_linestring_mask > 0
def _compute_box_mask(self, info, layers):
"""
Compute binary of bounding boxes in BEV space
:param annotations: annotation dataclass
:param layers: bounding box labels to include
:return: binary mask as numpy array
"""
box_polygon_mask = np.zeros(self._config.bev_semantic_frame[::-1], dtype=np.uint8)
for name_value, box_value in zip(info['anns']["gt_names"], info['anns']["gt_boxes"]):
agent_type = tracked_object_types[name_value]
if agent_type in layers:
# box_value = (x, y, z, length, width, height, yaw) TODO: add intenum
x, y, heading = box_value[0], box_value[1], box_value[-1]
box_length, box_width, box_height = box_value[3], box_value[4], box_value[5]
agent_box = OrientedBox(StateSE2(x, y, heading), box_length, box_width, box_height)
exterior = np.array(agent_box.geometry.exterior.coords).reshape((-1, 1, 2))
exterior = self._coords_to_pixel(exterior)
cv2.fillPoly(box_polygon_mask, [exterior], color=255)
# OpenCV has origin on top-left corner
box_polygon_mask = np.rot90(box_polygon_mask)[::-1]
return box_polygon_mask > 0
@staticmethod
def _geometry_local_coords(geometry, origin):
"""
Transform shapely geometry in local coordinates of origin.
:param geometry: shapely geometry
:param origin: pose dataclass
:return: shapely geometry
"""
a = np.cos(origin.heading)
b = np.sin(origin.heading)
d = -np.sin(origin.heading)
e = np.cos(origin.heading)
xoff = -origin.x
yoff = -origin.y
translated_geometry = affinity.affine_transform(geometry, [1, 0, 0, 1, xoff, yoff])
rotated_geometry = affinity.affine_transform(translated_geometry, [a, b, d, e, 0, 0])
return rotated_geometry
def _coords_to_pixel(self, coords):
"""
Transform local coordinates in pixel indices of BEV map
:param coords: _description_
:return: _description_
"""
# NOTE: remove half in backward direction
pixel_center = np.array([[0, self._config.bev_pixel_width / 2.0]])
coords_idcs = (coords / self._config.bev_pixel_size) + pixel_center
return coords_idcs.astype(np.int32) |