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)