scripts/analysis_tools/visualize/render/bev_render.py

import cv2
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from pyquaternion import Quaternion
from nuscenes.prediction import PredictHelper, convert_local_coords_to_global
from scripts.analysis_tools.visualize.render.base_render import BaseRender
from scripts.analysis_tools.visualize.utils import color_mapping, AgentPredictionData
from nuscenes.utils.data_classes import LidarPointCloud, Box
from mmdet3d_plugin.datasets.nuscenes_e2e_dataset import obtain_map_info
from scripts.analysis_tools.visualize.render.carla_map import get_carla_map_rasterize_semantic
from mmdet3d_plugin.datasets.nuplan.nuplan_pointcloud import PointCloud

class BEVRender(BaseRender):
    """
    Render class for BEV
    """

    def __init__(self,
                 figsize=(20, 20),
                 margin: float = 50,
                 view: np.ndarray = np.eye(4),
                 show_gt_boxes=False):
        super(BEVRender, self).__init__(figsize)
        self.margin = margin
        self.view = view
        self.show_gt_boxes = show_gt_boxes

    def set_plot_cfg(self):
        self.axes.set_xlim([-self.margin, self.margin])
        self.axes.set_ylim([-self.margin, self.margin])
        self.axes.set_aspect('equal')
        self.axes.grid(False)

    def render_sample_data(self, canvas, sample_token):
        pass

    def render_anno_data(
            self,
            sample_token,
            nusc,
            predict_helper):
        sample_record = nusc.get('sample', sample_token)
        assert 'LIDAR_TOP' in sample_record['data'].keys(
        ), 'Error: No LIDAR_TOP in data, unable to render.'
        lidar_record = sample_record['data']['LIDAR_TOP']
        data_path, boxes, _ = nusc.get_sample_data(
            lidar_record, selected_anntokens=sample_record['anns'])
        for box in boxes:
            instance_token = nusc.get('sample_annotation', box.token)[
                'instance_token']
            future_xy_local = predict_helper.get_future_for_agent(
                instance_token, sample_token, seconds=6, in_agent_frame=True)
            if future_xy_local.shape[0] > 0:
                trans = box.center
                rot = Quaternion(matrix=box.rotation_matrix)
                future_xy = convert_local_coords_to_global(
                    future_xy_local, trans, rot)
                future_xy = np.concatenate(
                    [trans[None, :2], future_xy], axis=0)
                c = np.array([0, 0.8, 0])
                box.render(self.axes, view=self.view, colors=(c, c, c))
                self._render_traj(future_xy, line_color=c, dot_color=(0, 0, 0))
        self.axes.set_xlim([-self.margin, self.margin])
        self.axes.set_ylim([-self.margin, self.margin])

    def show_lidar_data(
            self,
            sample_token,
            nusc,
            dataset,
            version,
            log_name=None,
            ):
        
        # retrieve lidar data
        if dataset == 'nusc':
            sample_record = nusc.get('sample', sample_token)
            assert 'LIDAR_TOP' in sample_record['data'].keys(
            ), 'Error: No LIDAR_TOP in data, unable to render.'
            lidar_record = sample_record['data']['LIDAR_TOP']
            data_path, boxes, _ = nusc.get_sample_data(
                lidar_record, selected_anntokens=sample_record['anns'])
            points = LidarPointCloud.from_file(data_path)
        elif dataset == 'carla':
            scene_token, frame_idstr = sample_token.split('_frame_')
            data_path = f'data/carla/{version}/val/{scene_token}/lidar_full/{frame_idstr}.npy'
            points = np.fromfile(data_path, dtype=np.float32)
            points = points.reshape(-1, 4)  # N x 4
            points = LidarPointCloud(points.T)
        elif dataset == 'nuplan':
            data_path = f'data/openscene-v1.1/sensor_blobs/{version}/{log_name}/MergedPointCloud/{sample_token}.pcd'
            points = PointCloud.parse_from_file(data_path).to_pcd_bin2()  # 6 x N
            points = points[:4].T  # N x 3
            points = LidarPointCloud(points.T)

        points.render_height(self.axes, view=self.view)
        self.axes.set_xlim([-self.margin, self.margin])
        self.axes.set_ylim([-self.margin, self.margin])
        self.axes.axis('off')
        self.axes.set_aspect('equal')

    def render_pred_box_data(self, agent_prediction_list):
        for pred_agent in agent_prediction_list:
            c = np.array([0, 1, 0])
            if hasattr(pred_agent, 'pred_track_id') and pred_agent.pred_track_id is not None:  # this is true
                tr_id = pred_agent.pred_track_id
                c = color_mapping[tr_id % len(color_mapping)]
            pred_agent.nusc_box.render(
                axis=self.axes, view=self.view, colors=(c, c, c))
            if pred_agent.is_sdc:
                c = np.array([1, 0, 0])
                pred_agent.nusc_box.render(
                    axis=self.axes, view=self.view, colors=(c, c, c))

    def render_pred_traj(self, agent_prediction_list, top_k=3):
        for pred_agent in agent_prediction_list:
            if pred_agent.is_sdc:
                continue
            sorted_ind = np.argsort(pred_agent.pred_traj_score)[
                ::-1]  # from high to low
            num_modes = len(sorted_ind)
            sorted_traj = pred_agent.pred_traj[sorted_ind, :, :2]
            sorted_score = pred_agent.pred_traj_score[sorted_ind]
            # norm_score = np.sum(np.exp(sorted_score))
            norm_score = np.exp(sorted_score[0])

            sorted_traj = np.concatenate(
                [np.zeros((num_modes, 1, 2)), sorted_traj], axis=1)
            trans = pred_agent.pred_center
            rot = Quaternion(axis=np.array([0, 0.0, 1.0]), angle=np.pi/2)
            vehicle_id_list = [0, 1, 2, 3, 4, 6, 7]
            if pred_agent.pred_label in vehicle_id_list:
                dot_size = 150
            else:
                dot_size = 25
            # print(sorted_score)
            for i in range(top_k-1, -1, -1):
                viz_traj = sorted_traj[i, :, :2]
                viz_traj = convert_local_coords_to_global(viz_traj, trans, rot)
                traj_score = np.exp(sorted_score[i])/norm_score
                # traj_score = [1.0, 0.01, 0.01, 0.01, 0.01, 0.01][i]
                self._render_traj(viz_traj, traj_score=traj_score,
                                  colormap='winter', dot_size=dot_size)

    def render_pred_map_data(self, predicted_map_seg, dataset):
        # rendered_map = map_color_dict
        # divider, crossing, drivable
        # orange, blue, green

        if predicted_map_seg is not None:
            map_color_dict = np.array(
                [(204, 128, 0), (102, 102, 255), (102, 255, 102)])
            rendered_map = map_color_dict[predicted_map_seg.argmax(
                -1).reshape(-1)].reshape(200, 200, -1)
            bg_mask = predicted_map_seg.sum(-1) == 0
            rendered_map[bg_mask, :] = 255      # H x W x 3
            
            # if dataset == 'carla':
            #     rendered_map = np.flip(rendered_map, axis=1)

            self.axes.imshow(rendered_map, alpha=0.6,
                            interpolation='nearest', extent=(-51.2, 51.2, -51.2, 51.2))

    def render_occ_map_data(self, agent_list):
        rendered_map = np.ones((200, 200, 3))
        rendered_map_hsv = matplotlib.colors.rgb_to_hsv(rendered_map)
        occ_prob_map = np.zeros((200, 200))
        for i in range(len(agent_list)):
            pred_agent = agent_list[i]
            if pred_agent.pred_occ_map is None:
                continue
            if hasattr(pred_agent, 'pred_track_id') and pred_agent.pred_track_id is not None:  # this is true
                tr_id = pred_agent.pred_track_id
                c = color_mapping[tr_id % len(color_mapping)]
            pred_occ_map = pred_agent.pred_occ_map.max(0)
            update_mask = pred_occ_map > occ_prob_map
            occ_prob_map[update_mask] = pred_occ_map[update_mask]
            pred_occ_map *= update_mask
            hsv_c = matplotlib.colors.rgb_to_hsv(c)
            rendered_map_hsv[pred_occ_map > 0.1] = (
                np.ones((200, 200, 1)) * hsv_c)[pred_occ_map > 0.1]
            max_prob = pred_occ_map.max()
            renorm_pred_occ_map = (pred_occ_map - max_prob) * 0.7 + 1
            sat_map = (renorm_pred_occ_map * hsv_c[1])
            rendered_map_hsv[pred_occ_map > 0.1,
                             1] = sat_map[pred_occ_map > 0.1]
            rendered_map = matplotlib.colors.hsv_to_rgb(rendered_map_hsv)
        self.axes.imshow(rendered_map, alpha=0.8,
                         interpolation='nearest', extent=(-50, 50, -50, 50))

    def render_occ_map_data_time(self, agent_list, t):
        rendered_map = np.ones((200, 200, 3))
        rendered_map_hsv = matplotlib.colors.rgb_to_hsv(rendered_map)
        occ_prob_map = np.zeros((200, 200))
        for i in range(len(agent_list)):
            pred_agent = agent_list[i]
            if pred_agent.pred_occ_map is None:
                continue
            if hasattr(pred_agent, 'pred_track_id') and pred_agent.pred_track_id is not None:  # this is true
                tr_id = pred_agent.pred_track_id
                c = color_mapping[tr_id % len(color_mapping)]
            pred_occ_map = pred_agent.pred_occ_map[t]
            update_mask = pred_occ_map > occ_prob_map
            occ_prob_map[update_mask] = pred_occ_map[update_mask]
            pred_occ_map *= update_mask
            hsv_c = matplotlib.colors.rgb_to_hsv(c)
            rendered_map_hsv[pred_occ_map > 0.1] = (
                np.ones((200, 200, 1)) * hsv_c)[pred_occ_map > 0.1]
            max_prob = pred_occ_map.max()
            renorm_pred_occ_map = (pred_occ_map - max_prob) * 0.7 + 1
            sat_map = (renorm_pred_occ_map * hsv_c[1])
            rendered_map_hsv[pred_occ_map > 0.1,
                             1] = sat_map[pred_occ_map > 0.1]
            rendered_map = matplotlib.colors.hsv_to_rgb(rendered_map_hsv)
        self.axes.imshow(rendered_map, alpha=0.8,
                         interpolation='nearest', extent=(-50, 50, -50, 50))

    def render_planning_data(self, predicted_planning, show_command=False, dataset=None):
        
        # render predicted trajectories
        planning_traj = predicted_planning.pred_traj    # 6 x 2
        planning_traj = np.concatenate(
            [np.zeros((1, 2)), planning_traj], axis=0)  # 7 x 2, add the current loaction (0,0)
        self._render_traj(planning_traj, colormap='autumn', dot_size=50)

        # render GT trajectories for ego
        planning_traj_gt = predicted_planning.traj_gt   # 6 x 2
        planning_traj_gt = np.concatenate(
            [np.zeros((1, 2)), planning_traj_gt], axis=0)
        self._render_traj(planning_traj_gt, colormap='cool', dot_size=50)

        # other color map options to be used
        # parula	Parula colormap array
        # turbo	Turbo colormap array (Since R2020b)
        # hsv	HSV colormap array
        # hot	Hot colormap array
        # cool	Cool colormap array
        # spring	Spring colormap array
        # summer	Summer colormap array
        # autumn	Autumn colormap array
        # winter	Winter colormap array
        # gray	Gray colormap array
        # bone	Bone colormap array
        # copper	Copper colormap array
        # pink	Pink colormap array
        # sky	Sky colormap array (Since R2023a)
        # abyss	Abyss colormap array (Since R2023b)
        # jet	Jet colormap array
        # lines	Lines colormap array
        # colorcube	Colorcube colormap array
        # prism	Prism colormap array
        # flag	Flag colormap array

        if show_command:
            self._render_command(predicted_planning.command, dataset)

    def render_planning_attn_mask(self, predicted_planning):
        planning_attn_mask = predicted_planning.attn_mask
        planning_attn_mask = planning_attn_mask/planning_attn_mask.max()
        cmap_name = 'plasma'
        self.axes.imshow(planning_attn_mask, alpha=0.8, interpolation='nearest', extent=(
            -51.2, 51.2, -51.2, 51.2), vmax=0.2, cmap=matplotlib.colormaps[cmap_name])

    def render_hd_map(self, nusc, map_data, sample_token, dataset, outputs):
        # import pdb;pdb.set_trace()

        if dataset == 'nusc':
            sample_record = nusc.get('sample', sample_token)
            sd_rec = nusc.get('sample_data', sample_record['data']['LIDAR_TOP'])
            cs_record = nusc.get('calibrated_sensor',
                                sd_rec['calibrated_sensor_token'])
            pose_record = nusc.get('ego_pose', sd_rec['ego_pose_token'])
            info = {
                'lidar2ego_translation': cs_record['translation'],
                'lidar2ego_rotation': cs_record['rotation'],
                'ego2global_translation': pose_record['translation'],
                'ego2global_rotation': pose_record['rotation'],
                'scene_token': sample_record['scene_token']
            }

            layer_names = ['road_divider', 'road_segment', 'lane_divider',
                        'lane',  'road_divider', 'traffic_light', 'ped_crossing']
            map_mask = obtain_map_info(nusc,
                                    map_data,
                                    info,
                                    patch_size=(102.4, 102.4),
                                    canvas_size=(1024, 1024),
                                    layer_names=layer_names)
            
        elif dataset == 'carla':
            map_mask = get_carla_map_rasterize_semantic(map_data, sample_token, outputs)    # 1 x H x W
            # map_mask = np.flip(map_mask, axis=1)

        # FLIP to match image
        map_mask = np.flip(map_mask, axis=1)
        map_mask = np.rot90(map_mask, k=-1, axes=(1, 2))
        
        # convert to binary and visualize
        map_mask = map_mask[:, ::-1] > 0
        # map_mask = map_mask > 0
        
        map_show = np.ones((1024, 1024, 3))
        map_show[map_mask[0], :] = np.array([1.00, 0.50, 0.31])
        if map_mask.shape[0] > 1:
            map_show[map_mask[1], :] = np.array([159./255., 0.0, 1.0])

        self.axes.imshow(map_show, alpha=0.2, interpolation='nearest',
                         extent=(-51.2, 51.2, -51.2, 51.2))

    def _render_traj(self, future_traj, traj_score=1, colormap='winter', points_per_step=20, line_color=None, dot_color=None, dot_size=25):
        total_steps = (len(future_traj)-1) * points_per_step + 1
        dot_colors = matplotlib.colormaps[colormap](
            np.linspace(0, 1, total_steps))[:, :3]
        dot_colors = dot_colors*traj_score + \
            (1-traj_score)*np.ones_like(dot_colors)
        total_xy = np.zeros((total_steps, 2))
        for i in range(total_steps-1):
            unit_vec = future_traj[i//points_per_step +
                                   1] - future_traj[i//points_per_step]
            total_xy[i] = (i/points_per_step - i//points_per_step) * \
                unit_vec + future_traj[i//points_per_step]
        total_xy[-1] = future_traj[-1]
        self.axes.scatter(
            total_xy[:, 0], total_xy[:, 1], c=dot_colors, s=dot_size)

    def _render_command(self, command, dataset):
        if dataset == 'nusc':
            # command_dict = ['TURN RIGHT', 'TURN LEFT', 'KEEP FORWARD']
            command_dict = ['TURN RIGHT', 'TURN LEFT', 'KEEP FORWARD', 'TURN RIGHT AT THE NEXT INTERSECTION', 'TURN LEFT AT THE NEXT INTERSECTION', 'PREPARE TO STOP ON THE LEFT', 'ENTER AND DRIVE IN THE ROUNDABOUT', 'EXIT THE ROUNDABOUT', 'UTURN']
        elif dataset == 'carla':
            command_dict = [
                "Turn Left",
                "Turn Right",
                "Go Straight",
                "Lane Follow",
                "Change Lane Left",
                "Change Lane Right",
            ]
        elif dataset == 'nuplan':
            command_dict = [
                "TURN LEFT",
                "KEEP FORWARD",
                "TURN RIGHT",
                "UNKNOWN",
            ]
        self.axes.text(-48, -45, command_dict[int(command)], fontsize=45)

    def render_sdc_car(self, dataset):
        sdc_car_png = cv2.imread('sources/sdc_car.png')
        sdc_car_png = cv2.cvtColor(sdc_car_png, cv2.COLOR_BGR2RGB)
        if dataset in ['nusc', 'carla']:
            self.axes.imshow(sdc_car_png, extent=(-1, 1, -2, 2))
        elif dataset == 'nuplan':
            sdc_car_png = cv2.rotate(sdc_car_png, cv2.ROTATE_90_CLOCKWISE)
            self.axes.imshow(sdc_car_png, extent=(-2, 2, -1, 1))

    def render_legend(self):
        legend = cv2.imread('sources/legend.png')
        legend = cv2.cvtColor(legend, cv2.COLOR_BGR2RGB)
        self.axes.imshow(legend, extent=(23, 51.2, -50, -40))