mmdet3d_plugin/datasets/data_utils/diffusiondrive_map.py

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)