| import numpy as np |
| from nuscenes.utils.data_classes import LidarPointCloud, Box |
| from pyquaternion import Quaternion |
|
|
|
|
| color_mapping = np.asarray([ |
| [0, 0, 0], |
| [255, 179, 0], |
| [128, 62, 117], |
| [255, 104, 0], |
| [166, 189, 215], |
| [193, 0, 32], |
| [206, 162, 98], |
| [129, 112, 102], |
| [0, 125, 52], |
| [246, 118, 142], |
| [0, 83, 138], |
| [255, 122, 92], |
| [83, 55, 122], |
| [255, 142, 0], |
| [179, 40, 81], |
| [244, 200, 0], |
| [127, 24, 13], |
| [147, 170, 0], |
| [89, 51, 21], |
| [241, 58, 19], |
| [35, 44, 22], |
| [112, 224, 255], |
| [70, 184, 160], |
| [153, 0, 255], |
| [71, 255, 0], |
| [255, 0, 163], |
| [255, 204, 0], |
| [0, 255, 235], |
| [255, 0, 235], |
| [255, 0, 122], |
| [255, 245, 0], |
| [10, 190, 212], |
| [214, 255, 0], |
| [0, 204, 255], |
| [20, 0, 255], |
| [255, 255, 0], |
| [0, 153, 255], |
| [0, 255, 204], |
| [41, 255, 0], |
| [173, 0, 255], |
| [0, 245, 255], |
| [71, 0, 255], |
| [0, 255, 184], |
| [0, 92, 255], |
| [184, 255, 0], |
| [255, 214, 0], |
| [25, 194, 194], |
| [92, 0, 255], |
| [220, 220, 220], |
| [255, 9, 92], |
| [112, 9, 255], |
| [8, 255, 214], |
| [255, 184, 6], |
| [10, 255, 71], |
| [255, 41, 10], |
| [7, 255, 255], |
| [224, 255, 8], |
| [102, 8, 255], |
| [255, 61, 6], |
| [255, 194, 7], |
| [0, 255, 20], |
| [255, 8, 41], |
| [255, 5, 153], |
| [6, 51, 255], |
| [235, 12, 255], |
| [160, 150, 20], |
| [0, 163, 255], |
| [140, 140, 140], |
| [250, 10, 15], |
| [20, 255, 0], |
| ])/255 |
|
|
|
|
| class AgentPredictionData: |
| """ |
| Agent data class, includes bbox, traj, and occflow |
| """ |
|
|
| def __init__(self, |
| pred_score, |
| pred_label, |
| pred_center, |
| pred_dim, |
| pred_yaw, |
| pred_vel, |
| pred_traj, |
| pred_traj_score, |
| pred_track_id=None, |
| pred_occ_map=None, |
| is_sdc=False, |
| past_pred_traj=None, |
| command=None, |
| attn_mask=None, |
| traj_gt=None, |
| ): |
| self.pred_score = pred_score |
| self.pred_label = pred_label |
| self.pred_center = pred_center |
| self.pred_dim = pred_dim |
| self.pred_vel = pred_vel |
| self.pred_traj = pred_traj |
| self.pred_traj_score = pred_traj_score |
| self.pred_track_id = pred_track_id |
| self.pred_occ_map = pred_occ_map |
| if self.pred_traj is not None: |
| if isinstance(self.pred_traj_score, int): |
| self.pred_traj_max = self.pred_traj |
| else: |
| self.pred_traj_max = self.pred_traj[self.pred_traj_score.argmax( |
| )] |
| else: |
| self.pred_traj_max = None |
|
|
| if pred_yaw is not None: |
| self.pred_yaw = -pred_yaw-np.pi/2 |
| self.nusc_box = Box( |
| center=pred_center, |
| size=pred_dim, |
| orientation=Quaternion(axis=[0, 0, 1], radians=self.pred_yaw), |
| label=pred_label, |
| score=pred_score |
| ) |
| else: |
| self.pred_yaw = None |
|
|
| if is_sdc: |
| self.pred_center = [0, 0, -1.2+1.56/2] |
| self.is_sdc = is_sdc |
| self.past_pred_traj = past_pred_traj |
| self.command = command |
| self.attn_mask = attn_mask |
| self.traj_gt = traj_gt |
