|
|
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 |
