scripts/data_converter/nuplan/e2e_nuplan_openscenes_navsim.py

import argparse
import glob
import os
import pickle
import pickle as pkl
from typing import Dict, List

import utils.sensor as box_sensor
import utils.transform as box_transform
import utils.vis as box_vis
import utils.rotation as object_rotation
import utils.calibration as calib_utils
import numpy as np
import pyquaternion
import scipy
from utils.miscell import find_unique_common_from_lists
from utils.data_structure import Box3D
from nuscenes.eval.common.utils import Quaternion, quaternion_yaw

from utils.nuplan_pointcloud import PointCloud

np.set_printoptions(precision=3, suppress=True)

FPS = 2
FRAME_INTERVAL = 1
FPS_KEYFRAME = FPS / FRAME_INTERVAL

global_track_id = 1
mapping_tracktoken2globalid = dict()


def parse_ego_sensor_calib(anno: Dict, data_split: str, vis: bool = False) -> Dict:
    """Parse calibration between world, ego, lidar and camera"""

    # CAN_BUS definition
    # https://github.com/OpenDriveLab/OpenScene/blob/main/DriveEngine/process_data/helpers/canbus.py
    # self.tensor = np.array(
    #     [
    #         self.x,
    #         self.y,
    #         self.z,
    #         self.qw,
    #         self.qx,
    #         self.qy,
    #         self.qz,
    #         self.acceleration_x,
    #         self.acceleration_y,
    #         self.acceleration_z,
    #         self.vx,
    #         self.vy,
    #         self.vz,
    #         self.angular_rate_x,
    #         self.angular_rate_y,
    #         self.angular_rate_z,
    #         0.0,
    #         0.0,
    #     ]
    # )

    # sdc's box in lidar coordinate
    sdc_center_ego = np.array([[0, 0, 0, 1.0]]).transpose()  # 4 x 1
    sdc_center_lidar = np.linalg.inv(anno["lidar2ego"]) @ sdc_center_ego  # 4 x 1
    yaw_lidar = scipy.spatial.transform.Rotation.from_matrix(
        np.linalg.inv(anno["lidar2ego"])[:3, :3]
    ).as_euler("xyz", degrees=False)[-1]
    l, w, h = 4.52, 2.1, 1.56
    sdc_vel_ego = anno["can_bus"][10:13]
    sdc_vel_global = anno["ego2global"][:3, :3] @ sdc_vel_ego
    sdc_vel_lidar = np.linalg.inv(anno["lidar2ego"][:3, :3]) @ sdc_vel_ego
    # print('sdc_vel_global', sdc_vel_global)
    # print(sdc_vel_ego.shape)
    # print(sdc_vel_lidar.shape)

    gt_sdc_bbox_lidar = np.array(
        [
            sdc_center_lidar[0, 0],
            sdc_center_lidar[1, 0],
            sdc_center_lidar[2, 0],
            l,
            w,
            h,
            yaw_lidar,
            sdc_vel_lidar[0],
            sdc_vel_lidar[1],
        ]
    )  # 9
    # print('scene_token', anno['scene_token'])
    # print('frame_idx', anno['frame_idx'])

    # load camera calib
    cams = anno["cams"]
    for cam_type, cam_dict in cams.items():
        sensor2lidar = np.identity(4)
        sensor2lidar[:3, :3] = cam_dict["sensor2lidar_rotation"]
        sensor2lidar[:3, 3] = cam_dict["sensor2lidar_translation"]
        sensor2ego: np.ndarray = anno["lidar2ego"] @ sensor2lidar

        # construct nuScenes-like camera record
        cams[cam_type].update(
            {
                "type": cam_type,
                "sensor2ego_translation": sensor2ego[:3, 3],
                "sensor2ego_rotation": object_rotation.convert_mat2qua(
                    sensor2ego[:3, :3]
                ),
            }
        )

    pts_filename = anno["lidar_path"]  # log_name + token
    src_split = data_split.split("_")[0]
    data_root = f"./data/openscene-v1.1/sensor_blobs/{src_split}"
    if data_root not in pts_filename:
        pts_filename = os.path.join(data_root, pts_filename)

    points = PointCloud.parse_from_file(pts_filename).to_pcd_bin2()  # 6 x N
    points = points[:3].T  # N x 3

    # visualize lidar points only, default color ORANGE
    if vis:
        log_root = pts_filename.split("/")[5]
        save_root = (
            f"./data/openscene-v1.1/vis/{data_split}/{log_root}/{anno['scene_name']}"
        )

        if not os.path.exists(save_root):
            os.makedirs(save_root)
        save_path = os.path.join(
            save_root, f"{anno['frame_idx']:06d}_{anno['token']}.jpg"
        )
        box_vis.vis_box_on_lidar(
            lidar=points,
            save_path=save_path,
            left_hand=False,
            ego2lidar=np.linalg.inv(anno["lidar2ego"]),
        )

    command = np.argmax(anno["driving_command"])  # int
    # update dictionary
    anno.update(
        {
            "lidar_pts": points,
            "sweeps": [],
            "cams": cams,
            "sdc_vel_global": sdc_vel_global,
            # sdc's information in lidar coordinate
            "gt_sdc_bbox_lidar": gt_sdc_bbox_lidar,  # 9
            "command": command,  # int, convert 1-indexed -> 0-indexed
        }
    )

    return anno


def parse_bbox(anno: Dict, data_split: str, vis: bool = False) -> Dict:
    """Parse each one of the bounding boxes and compute properties"""

    num_obj: int = anno["anns"]["gt_boxes"].shape[0]
    anno.update(
        {
            "gt_velocity": anno["anns"]["gt_velocity_3d"][:, :2],  # lidar coordinate
            "gt_boxes": anno["anns"]["gt_boxes"],  # lidar coordinate (x,y,z,l,w,h,yaw)
            "gt_bboxes_global": np.zeros((num_obj, 9), dtype="float32"),
            "gt_names": anno["anns"]["gt_names"],
            "gt_inds": np.zeros((num_obj), dtype="int64"),
            "num_lidar_pts": np.zeros((num_obj,), dtype="float32"),
            "valid_flag": np.zeros((num_obj,), dtype="bool"),
        }
    )

    # go through all bbox to fill in database info
    bbox_gt_in_ego_list: List[np.ndarray] = list()
    valid_id_list: List[np.ndarray] = list()
    for bbox_index in range(num_obj):
        # retrieve info
        box_7dof_lidar = anno["anns"]["gt_boxes"][bbox_index]
        instance_token = anno["anns"]["instance_tokens"][bbox_index]
        track_token = anno["anns"]["track_tokens"][bbox_index]

        # assign global ID
        global mapping_tracktoken2globalid
        if track_token not in mapping_tracktoken2globalid:
            global global_track_id
            mapping_tracktoken2globalid[track_token] = global_track_id
            anno["gt_inds"][bbox_index] = global_track_id

            # update track_id for the next object
            global_track_id += 1
        else:
            anno["gt_inds"][bbox_index] = mapping_tracktoken2globalid[track_token]

        # TODO check how they change over frame
        # print('instance_token', instance_token)
        # print('track_token', track_token)
        # print('track_id', anno["gt_inds"][bbox_index])
        # zxc

        # get box in global coordinate
        bbox2lidar = calib_utils.create_transform(
            box_7dof_lidar[:3], np.array([0, box_7dof_lidar[-1], 0])
        )  # 4 x 4
        # if anno["gt_inds"][bbox_index] == 19:
        # print('bbox2lidar\n', bbox2lidar)

        # convert from lidar to world
        bbox_center_lidar = np.array([[0, 0, 0, 1.0]]).transpose()  # 4 x 1
        bbox_center_lidar[:3, 0] = box_7dof_lidar[:3]
        l, w, h = box_7dof_lidar[3:6]
        bbox_center_world = anno["lidar2global"] @ bbox_center_lidar
        bbox2global = anno["lidar2global"] @ bbox2lidar
        yaw_global = scipy.spatial.transform.Rotation.from_matrix(
            bbox2global[:3, :3]
        ).as_euler("xyz", degrees=False)[-1]
        velocity_lidar = np.array([[0, 0, 1.0]]).transpose()  # 3 x 1
        velocity_lidar[:2, 0] = anno["gt_velocity"][bbox_index]  # 3 x 1
        velocity_world = anno["lidar2global"][:3, :3] @ velocity_lidar

        # get 7dof representation: xyzlwh + yaw, global coordinate
        box_7dof_global = np.array(
            [
                float(bbox_center_world[0]),
                float(bbox_center_world[1]),
                float(bbox_center_world[2]),
            ]
            + [l, w, h]
            + [float(yaw_global)]
        )  # (7, )
        anno["gt_bboxes_global"][bbox_index, :7] = box_7dof_global
        anno["gt_bboxes_global"][bbox_index, 7:] = velocity_world[:2, 0]
        # if anno["gt_inds"][bbox_index] == 19:
        # print(anno["gt_bboxes_global"][bbox_index])

        # extract number of points
        box3d: Box3D = Box3D.array2bbox_xyzlwhyaw(box_7dof_lidar)
        box_in_lidar: np.ndarray = Box3D.box2corners3d_lidar(box3d)  # 8 x 4
        pc_sub: np.ndarray = box_sensor.extract_pc_in_box3d(
            anno["lidar_pts"], box_in_lidar
        )[
            0
        ]  # N x 3
        # print('number of points wihin the point cloud is %d' % pc_sub.shape[0])

        # add data into list
        anno["num_lidar_pts"][bbox_index] = pc_sub.shape[0]

        # set valid only if there are lidar points inside the box
        if pc_sub.shape[0] > 0:
            anno["valid_flag"][bbox_index] = True

            # convert to the ego coordinate for visualization
            if vis:
                box_in_ego: np.ndarray = box_transform.box_in_lidar_to_ego(
                    box_in_lidar, lidar2ego=anno["lidar2ego"],
                )  # 8 x 4
                bbox_gt_in_ego_list.append(box_in_ego)
                valid_id_list.append(anno["gt_inds"][bbox_index])

    if vis:
        pts_filename = anno["lidar_path"]  # log_name + token
        log_root = pts_filename.split("/")[0]
        save_root = f"./data/openscene-v1.1/vis_bbox/{data_split}/{log_root}/{anno['scene_name']}"

        if not os.path.exists(save_root):
            os.makedirs(save_root)
        save_path = os.path.join(
            save_root, f"{anno['frame_idx']:06d}_{anno['token']}.jpg"
        )

        # add ego box
        box3d: Box3D = Box3D.array2bbox_xyzlwhyaw(anno["gt_sdc_bbox_lidar"][:7])
        sdc_box_in_lidar: np.ndarray = Box3D.box2corners3d_lidar(box3d)  # 8 x 4
        sdc_box_in_ego: np.ndarray = box_transform.box_in_lidar_to_ego(
            sdc_box_in_lidar, lidar2ego=anno["lidar2ego"],
        )  # 8 x 4
        bbox_gt_in_ego_list.append(sdc_box_in_ego)
        valid_id_list.append(-1)

        # visualizating boxes
        box_vis.vis_box_on_lidar(
            anno["lidar_pts"],
            save_path,
            bbox_gt_in_ego_list=bbox_gt_in_ego_list,
            bev=True,
            id_list=valid_id_list,
            left_hand=False,
            ego2lidar=np.linalg.inv(anno["lidar2ego"]),
        )

    return anno


def parse_data(anno: Dict, data_split: str, vis: bool = False) -> Dict:
    """Convert carla data into the collection of path needed by mmdetection3D dataloader"""

    anno: Dict = parse_ego_sensor_calib(anno, data_split, vis=vis)
    anno: Dict = parse_bbox(anno, data_split, vis=vis)

    return anno


def parse_sequence(
    seq: str,
    annos: List[Dict],
    seq_index: List[int],
    data_split: str,
    pre_sec: float = 1.5,
    fut_sec: float = 4,
    vis: bool = False,
) -> List[Dict]:
    """Add past/future trajectory into data at each frame to allow prediction"""

    anno_seq: List[Dict] = [annos[index] for index in seq_index]
    print(f"processing: {seq}, number of frames {len(anno_seq)}")

    # sort data by frame ID, as the original data might not be sorted based
    # on frame index
    def sort_data(data_dict: Dict):
        return data_dict["frame_idx"]
    anno_seq.sort(key=sort_data)

    # check if the sequence has non-continuous frames, if so, remove the data
    min_frame_ID: int = anno_seq[0]["frame_idx"]
    max_frame_ID: int = anno_seq[-1]["frame_idx"]
    if len(anno_seq) < max_frame_ID - min_frame_ID + 1:
        for index in seq_index:
            annos[index] = None
        
        return annos

    #     # return 

    # print('min_frame_ID', min_frame_ID)
    # print('max_frame_ID', max_frame_ID)
    # print('length', len(anno_seq))

    # loop through each frame in this sequence
    for index in range(len(seq_index)):
        anno: Dict = anno_seq[index]  # data for the frame
        frame_ID: int = anno["frame_idx"]  # the frame ID
        # print("current frame is", frame_ID)

        # check if the data is from the same sequence
        scene_token: str = anno["scene_token"]  # the sequence name
        assert scene_token == seq, "error, not the same sequence"

        # retrieve info for other objects
        # gt_velocity: np.ndarray = anno["gt_velocity"]  # N x 2
        gt_bboxes_lidar: np.ndarray = anno["anns"]["gt_boxes"]  # N x 7
        num_obj: int = gt_bboxes_lidar.shape[0]

        # convert the str ID to global ID
        obj_ids: List[int] = anno["gt_inds"].tolist()
        # print("obj_ids is", obj_ids)
        # print("num_obj is", num_obj)

        # retrieve info for ego and calib
        l, w, h = anno["gt_sdc_bbox_lidar"][3:6]
        lidar2global: np.ndarray = anno["lidar2global"]  # 4 x 4

        ########### get the target data we need for GT
        pre_frames: int = int(pre_sec * FPS_KEYFRAME)
        fut_frames: int = int(fut_sec * FPS_KEYFRAME)

        # the past & future trajectories for objects existing in the current frame
        # so we know the exact number of objects
        gt_fut_bbox_lidar: np.ndarray = np.zeros((num_obj, fut_frames, 9))
        gt_pre_bbox_lidar: np.ndarray = np.zeros((num_obj, pre_frames, 9))

        # initially set all mask as invalid until we identify valid frames
        gt_fut_bbox_mask: np.ndarray = np.zeros((num_obj, fut_frames, 1))
        gt_pre_bbox_mask: np.ndarray = np.zeros((num_obj, pre_frames, 1))

        # ALL the past & future trajectories, including objects that appear
        # from frames to frames, and not existing in the current frame
        # as a result, the number of objects might not be consistent over frames
        # can be used to create an Union of all objects for better training
        gt_fut_bbox_lidar_all: List[np.ndarray] = [
            [np.empty([0])] for _ in range(fut_frames)
        ]
        gt_pre_bbox_lidar_all: List[np.ndarray] = [
            [np.empty([0])] for _ in range(pre_frames)
        ]

        # sdc's temporal info
        gt_pre_bbox_sdc_lidar: np.ndarray = np.zeros((1, pre_frames, 9))
        gt_fut_bbox_sdc_lidar: np.ndarray = np.zeros((1, fut_frames, 9))
        gt_pre_bbox_sdc_global: np.ndarray = np.zeros((1, pre_frames, 9))
        gt_fut_bbox_sdc_global: np.ndarray = np.zeros((1, fut_frames, 9))
        gt_pre_bbox_sdc_mask: np.ndarray = np.zeros((1, pre_frames, 1))
        gt_fut_bbox_sdc_mask: np.ndarray = np.zeros((1, fut_frames, 1))
        gt_pre_command_sdc: np.ndarray = np.zeros((1, pre_frames, 1))
        gt_fut_command_sdc: np.ndarray = np.zeros((1, fut_frames, 1))

        # get past trajectory, backward order, i.e., frame 4, 3, 2, 1
        # start with 1 to not include the current frame
        for pre_frame_index in range(1, pre_frames + 1):
            pre_frame_ID = int(frame_ID - pre_frame_index * FRAME_INTERVAL)
            pre_frame_index_in_seq: int = index - pre_frame_index

            # skip the initial frame with incomplete lidar and also negative frame
            if pre_frame_ID < min_frame_ID:
                continue

            # retrieve the data in the global coordinate
            anno_pre_tmp: Dict = anno_seq[pre_frame_index_in_seq]
            gt_pre_bbox_global_tmp: np.ndarray = anno_pre_tmp[
                "gt_bboxes_global"
            ]  # N x 9
            num_obj_pre: int = gt_pre_bbox_global_tmp.shape[0]
            gt_pre_bbox_lidar_tmp = np.zeros((num_obj_pre, 9))

            ########## convert the global coordinate to lidar coordinate in current frame
            # i.e., ego motion compensation

            # location
            loc_global = np.concatenate(
                (gt_pre_bbox_global_tmp[:, :3], np.ones((num_obj_pre, 1))), axis=1
            ).T  # 4 x N
            loc_lidar = np.linalg.inv(lidar2global) @ loc_global  # 4 x N
            gt_pre_bbox_lidar_tmp[:, :3] = loc_lidar[:3, :].transpose()  # N x 3

            # size
            gt_pre_bbox_lidar_tmp[:, 3:6] = gt_pre_bbox_global_tmp[:, 3:6]

            # rotation
            gt_pre_rot_global_rad = np.concatenate(
                (
                    np.zeros((num_obj_pre, 1)),
                    gt_pre_bbox_global_tmp[:, [6]],
                    np.zeros((num_obj_pre, 1)),
                ),
                axis=1,
            )  # N x 3
            gt_pre_rot_global_deg = gt_pre_rot_global_rad / np.pi * 180
            for obj_index in range(num_obj_pre):
                gt_pre_trans_global = calib_utils.create_transform(
                    gt_pre_bbox_global_tmp[obj_index, :3],
                    gt_pre_rot_global_deg[obj_index],
                )  # 4 x 4
                gt_pre_trans_lidar = (
                    np.linalg.inv(anno["lidar2global"]) @ gt_pre_trans_global
                )  # 4 x 4
                gt_pre_yaw_lidar = scipy.spatial.transform.Rotation.from_matrix(
                    gt_pre_trans_lidar[:3, :3]
                ).as_euler("xyz", degrees=False)[-1]
                gt_pre_bbox_lidar_tmp[obj_index, 6] = gt_pre_yaw_lidar  # N x 3

            # velocity
            gt_pre_vel_global = np.concatenate(
                (gt_pre_bbox_global_tmp[:, 7:], np.zeros((num_obj_pre, 1)),),  # N x 2
                axis=1,
            )  # N x 3
            gt_pre_vel_lidar: np.ndarray = (
                np.linalg.inv(lidar2global[:3, :3]) @ gt_pre_vel_global.T
            ).T  # N x 3
            gt_pre_bbox_lidar_tmp[:, 7:] = gt_pre_vel_lidar[:, :2]  # N x 2

            # dump the data into the set of all objects in the future
            gt_pre_bbox_lidar_all[pre_frame_index - 1] = gt_pre_bbox_lidar_tmp  # N x 9

            # now check the IDs that exist in the current frame
            # in order to produce the mask for future frames
            obj_ids_pre: List[int] = anno_pre_tmp["gt_inds"].tolist()
            (
                obj_ids_common,
                index_cur,
                index_past,
            ) = find_unique_common_from_lists(obj_ids, obj_ids_pre)

            # possible that we get 0 object in future frame matching with GT object currint
            if len(index_cur) > 0:
                gt_pre_bbox_mask[np.array(index_cur), pre_frame_index - 1] = 1

                # also assign the actual gt boxes into the array
                gt_pre_bbox_lidar[
                    np.array(index_cur), pre_frame_index - 1, :
                ] = gt_pre_bbox_lidar_tmp[np.array(index_past), :]

            ############### get past of the sdc box in lidar coordinate
            sdc_loc_global, _ = calib_utils.transform_matrix_to_vector(
                anno_pre_tmp["ego2global"]
            )
            sdc_loc_global = np.concatenate((sdc_loc_global, np.array([1])))  # 4
            sdc_loc_lidar = np.linalg.inv(lidar2global) @ sdc_loc_global  # 4
            sdc_vel_global: np.ndarray = anno_pre_tmp["sdc_vel_global"]  # 3
            sdc_vel_lidar: np.ndarray = (
                np.linalg.inv(lidar2global[:3, :3]) @ sdc_vel_global
            )  # 3
            sdc_trans_lidar = (
                np.linalg.inv(lidar2global) @ anno_pre_tmp["ego2global"]
            )  # 4 x 4
            yaw_lidar = scipy.spatial.transform.Rotation.from_matrix(
                sdc_trans_lidar[:3, :3]
            ).as_euler("xyz", degrees=False)[-1]
            yaw_global = scipy.spatial.transform.Rotation.from_matrix(
                anno_pre_tmp["ego2global"][:3, :3]
            ).as_euler("xyz", degrees=False)[-1]
            gt_pre_bbox_sdc_mask[0, pre_frame_index - 1] = 1

            # put things into bbox, 9 dof
            gt_sdc_bbox_lidar = np.array(
                [
                    sdc_loc_lidar[0],
                    sdc_loc_lidar[1],
                    sdc_loc_lidar[2],
                    l,
                    w,
                    h,
                    yaw_lidar,
                    sdc_vel_lidar[0],
                    sdc_vel_lidar[1],
                ]
            )
            gt_pre_bbox_sdc_lidar[0, pre_frame_index - 1] = gt_sdc_bbox_lidar
            gt_sdc_bbox_global = np.array(
                [
                    sdc_loc_global[0],
                    sdc_loc_global[1],
                    sdc_loc_global[2],
                    l,
                    w,
                    h,
                    yaw_global,
                    sdc_vel_global[0],
                    sdc_vel_global[1],
                ]
            )
            gt_pre_bbox_sdc_global[0, pre_frame_index - 1] = gt_sdc_bbox_global

            # get command
            gt_pre_command_sdc[0, pre_frame_index - 1]: int = anno_pre_tmp["command"]

        # get future trajectory
        # start with 1 to not include the current frame
        for fut_frame_index in range(1, fut_frames + 1):
            fut_frame_ID = int(frame_ID + fut_frame_index * FRAME_INTERVAL)
            fut_frame_index_in_seq: int = index + fut_frame_index

            # beyond the last timestamp
            if fut_frame_ID > max_frame_ID:
                continue

            # retrieve the data in the global coordinate
            try:
                anno_fut_tmp: Dict = anno_seq[fut_frame_index_in_seq]
            except:
                print('fut_frame_index_in_seq', fut_frame_index_in_seq)
                anno_fut_tmp: Dict = anno_seq[fut_frame_index_in_seq]

            gt_fut_bbox_global_tmp: np.ndarray = anno_fut_tmp[
                "gt_bboxes_global"
            ]  # N x 9
            num_obj_fut: int = gt_fut_bbox_global_tmp.shape[0]
            gt_fut_bbox_lidar_tmp = np.zeros((num_obj_fut, 9))

            ########## convert the global coordinate to lidar coordinate in current frame
            # i.e., ego motion compensation

            # location
            loc_global = np.concatenate(
                (gt_fut_bbox_global_tmp[:, :3], np.ones((num_obj_fut, 1))), axis=1,
            ).T  # 4 x N
            loc_lidar = np.linalg.inv(lidar2global) @ loc_global  # 4 x N
            gt_fut_bbox_lidar_tmp[:, :3] = loc_lidar[:3, :].transpose()  # N x 3

            # size
            gt_fut_bbox_lidar_tmp[:, 3:6] = gt_fut_bbox_global_tmp[:, 3:6]

            # rotation
            gt_fut_rot_global_rad = np.concatenate(
                (
                    np.zeros((num_obj_fut, 1)),
                    gt_fut_bbox_global_tmp[:, [6]],
                    np.zeros((num_obj_fut, 1)),
                ),
                axis=1,
            )  # N x 3
            gt_fut_rot_global_deg = gt_fut_rot_global_rad / np.pi * 180
            for obj_index in range(num_obj_fut):
                gt_trans_global_future = calib_utils.create_transform(
                    gt_fut_bbox_global_tmp[obj_index, :3],
                    gt_fut_rot_global_deg[obj_index],
                )  # 4 x 4
                gt_fut_trans_lidar = (
                    np.linalg.inv(anno["lidar2global"]) @ gt_trans_global_future
                )  # 4 x 4
                gt_fut_yaw_lidar = scipy.spatial.transform.Rotation.from_matrix(
                    gt_fut_trans_lidar[:3, :3]
                ).as_euler("xyz", degrees=False)[-1]
                gt_fut_bbox_lidar_tmp[obj_index, 6] = gt_fut_yaw_lidar  # N x 3

            # velocity
            gt_fut_vel_global = np.concatenate(
                (gt_fut_bbox_global_tmp[:, 7:], np.zeros((num_obj_fut, 1)),),  # N x 2
                axis=1,
            )  # N x 3
            gt_fut_vel_lidar: np.ndarray = (
                np.linalg.inv(lidar2global[:3, :3]) @ gt_fut_vel_global.T
            ).T  # N x 3
            gt_fut_bbox_lidar_tmp[:, 7:] = gt_fut_vel_lidar[:, :2]  # N x 2

            # gt_vel_future_tmp: np.ndarray = anno_future_tmp["gt_velocity"]  # N x 2
            # gt_box_data_all: np.ndarray = np.concatenate(
            #     (gt_boxes_future_tmp, gt_vel_future_tmp), axis=-1
            # )  # N x 9
            # print("\nfuture frame is", frame_idx)
            # print(gt_boxes_future_tmp)

            # dump the data into the set of all objects in the future
            gt_fut_bbox_lidar_all[fut_frame_index - 1] = gt_fut_bbox_lidar_tmp  # N x 9

            # now check the IDs that exist in the current frame
            # in order to produce the mask for future frames
            obj_ids_fut: List[int] = anno_fut_tmp["gt_inds"].tolist()
            (
                obj_ids_common,
                index_cur,
                index_fut,
            ) = find_unique_common_from_lists(obj_ids, obj_ids_fut)

            # possible that we get 0 object in future frame matching with GT object currint
            if len(index_cur) > 0:
                gt_fut_bbox_mask[np.array(index_cur), fut_frame_index - 1] = 1
                # also assign the actual gt boxes into the array
                gt_fut_bbox_lidar[
                    np.array(index_cur), fut_frame_index - 1, :
                ] = gt_fut_bbox_lidar_tmp[np.array(index_fut), :]

            ############## get future of the sdc traj
            # gt_fut_sdc_global: np.ndarray = anno["can_bus"][:3]
            # gt_fut_sdc_global = np.concatenate(
            # (gt_fut_sdc_global, np.array([1]))
            # )
            # gt_fut_sdc_lidar = (
            # np.linalg.inv(anno["lidar2global"]) @ gt_fut_sdc_global
            # )  # 4 x 1
            # gt_fut_traj_sdc_lidar[0, future_frame_index - 1, :] = gt_fut_sdc_lidar[:3]

            ############### get future of the sdc box in lidar coordinate
            sdc_loc_global, _ = calib_utils.transform_matrix_to_vector(
                anno_fut_tmp["ego2global"]
            )
            sdc_loc_global = np.concatenate((sdc_loc_global, np.array([1])))  # 4
            sdc_loc_lidar = np.linalg.inv(lidar2global) @ sdc_loc_global  # 4
            sdc_vel_global: np.ndarray = anno_fut_tmp["sdc_vel_global"]  # 3
            sdc_vel_lidar: np.ndarray = (
                np.linalg.inv(lidar2global[:3, :3]) @ sdc_vel_global
            )  # 3
            sdc_trans_lidar = (
                np.linalg.inv(lidar2global) @ anno_fut_tmp["ego2global"]
            )  # 4 x 4
            yaw_lidar = scipy.spatial.transform.Rotation.from_matrix(
                sdc_trans_lidar[:3, :3]
            ).as_euler("xyz", degrees=False)[-1]
            yaw_global = scipy.spatial.transform.Rotation.from_matrix(
                anno_fut_tmp["ego2global"][:3, :3]
            ).as_euler("xyz", degrees=False)[-1]
            gt_fut_bbox_sdc_mask[0, fut_frame_index - 1] = 1

            # put things into bbox, 9 dof
            gt_sdc_bbox_lidar = np.array(
                [
                    sdc_loc_lidar[0],
                    sdc_loc_lidar[1],
                    sdc_loc_lidar[2],
                    l,
                    w,
                    h,
                    yaw_lidar,
                    sdc_vel_lidar[0],
                    sdc_vel_lidar[1],
                ]
            )
            gt_fut_bbox_sdc_lidar[0, fut_frame_index - 1] = gt_sdc_bbox_lidar
            gt_sdc_bbox_global = np.array(
                [
                    sdc_loc_global[0],
                    sdc_loc_global[1],
                    sdc_loc_global[2],
                    l,
                    w,
                    h,
                    yaw_global,
                    sdc_vel_global[0],
                    sdc_vel_global[1],
                ]
            )
            gt_fut_bbox_sdc_global[0, fut_frame_index - 1] = gt_sdc_bbox_global

            # get command
            gt_fut_command_sdc[0, fut_frame_index - 1]: int = anno_fut_tmp["command"]

        # flip array to allow time in forward order for the past trajectory
        gt_pre_bbox_lidar = np.flip(gt_pre_bbox_lidar, axis=1)
        gt_pre_bbox_mask = np.flip(gt_pre_bbox_mask, axis=1)
        gt_pre_bbox_sdc_lidar = np.flip(gt_pre_bbox_sdc_lidar, axis=1)
        gt_pre_bbox_sdc_global = np.flip(gt_pre_bbox_sdc_global, axis=1)
        gt_pre_bbox_sdc_mask = np.flip(gt_pre_bbox_sdc_mask, axis=1)
        gt_pre_command_sdc = np.flip(gt_pre_command_sdc, axis=1)

        # assign the value back into the dictionary
        # in-place assignment to the sorted anno
        anno["gt_pre_bbox_lidar"] = gt_pre_bbox_lidar  # N x 4 x 9
        anno["gt_fut_bbox_lidar"] = gt_fut_bbox_lidar  # N x 12 x 9
        anno["gt_pre_bbox_mask"] = gt_pre_bbox_mask  # N x 4 x 1
        anno["gt_fut_bbox_mask"] = gt_fut_bbox_mask  # N x 12 x 1
        anno["gt_pre_bbox_sdc_lidar"] = gt_pre_bbox_sdc_lidar  # 1 x 4 x 9
        anno["gt_fut_bbox_sdc_lidar"] = gt_fut_bbox_sdc_lidar  # 1 x 12 x 9
        anno["gt_pre_bbox_sdc_global"] = gt_pre_bbox_sdc_global  # 1 x 4 x 9
        anno["gt_fut_bbox_sdc_global"] = gt_fut_bbox_sdc_global  # 1 x 12 x 9
        anno["gt_pre_bbox_sdc_mask"] = gt_pre_bbox_sdc_mask  # 1 x 4 x 1
        anno["gt_fut_bbox_sdc_mask"] = gt_fut_bbox_sdc_mask  # 1 x 12 x 1
        anno["gt_pre_command_sdc"] = gt_pre_command_sdc  # 1 x 4 x 1
        anno["gt_fut_command_sdc"] = gt_fut_command_sdc  # 1 x 12 x 1

        # debug other object's trajectory
        # print('\nobjects')
        # # index = 6
        # # obj_id = obj_ids[index]
        # # print(obj_id)

        # obj_id = 19
        # index = obj_ids.index(obj_id)

        # print(index)
        # gt_bboxes_lidar = np.expand_dims(gt_bboxes_lidar, axis=1)
        # gt_box_curfut = np.concatenate(
        #     (gt_bboxes_lidar, gt_fut_traj), axis=1
        # )  # N x 13 x 9
        # print('gt_pre_bbox_lidar\n', gt_pre_bbox_lidar[index])
        # print('gt_fut_bbox_lidar\n', gt_fut_bbox_lidar[index])
        # print('gt_pre_bbox_mask\n', gt_pre_bbox_mask[index])
        # print('gt_fut_bbox_mask\n', gt_fut_bbox_mask[index])
        # print('gt_pre_bbox_lidar\n', gt_pre_bbox_lidar.shape)
        # print('gt_fut_bbox_lidar\n', gt_fut_bbox_lidar.shape)
        # print('gt_pre_bbox_mask\n', gt_pre_bbox_mask.shape)
        # print('gt_fut_bbox_mask\n', gt_fut_bbox_mask.shape)
        # zxc

        # debug ego trajectory
        # gt_sdc_bbox_lidar_all = np.concatenate(
        #     (
        #         gt_pre_bbox_sdc_lidar[0],  # 4 x 9
        #         anno["gt_sdc_bbox_lidar"].reshape((1, -1)),  # 1 x 9
        #         gt_fut_bbox_sdc_lidar[0],  # 12 x 9
        #     ),
        #     axis=0,
        # )
        # gt_sdc_mask_all = np.concatenate(
        #     (
        #         gt_pre_bbox_sdc_mask,  # 1 x 4 x 1
        #         np.array([1]).reshape(1, -1, 1),
        #         gt_fut_bbox_sdc_mask,
        #     ),
        #     axis=1,
        # )
        # gt_sdc_command_all = np.concatenate(
        #     (
        #         gt_pre_command_sdc,
        #         np.array([anno["command"]]).reshape(1, -1, 1),
        #         gt_fut_command_sdc),
        #     axis=1,
        # )
        # print(gt_fut_traj_sdc_lidar)
        # print("\nego")
        # print('gt_sdc_bbox_lidar_all\n', gt_sdc_bbox_lidar_all)
        # print(gt_sdc_mask_all)
        # print(gt_sdc_command_all)
        # zxc

        # data_frame = [anno['frame_idx'] for anno in annos]
        # print(data_frame)
        # data_frame = [anno['gt_fut_traj'].shape for anno in annos]
        # print(data_frame)

        #         print(obj_ids)
        #         print(obj_ids_fut)
        #         print(obj_ids_common)
        #         print(index_cur)
        #         print(index_fut)
        #         print(gt_fut_traj_mask[:, future_frame_index - 1])
        #         print(gt_fut_traj[:, future_frame_index - 1])
        #         # zxc
        #     zxc
        # zxc

        if vis:
            # get save path for train/val data_split

            pts_filename = anno["lidar_path"]  # log_name + token
            log_root = pts_filename.split("/")[0]
            save_root = f"./data/openscene-v1.1/vis_seq/{data_split}/{log_root}/{anno['scene_name']}"
            if not os.path.exists(save_root):
                os.makedirs(save_root)
            save_path = os.path.join(
                save_root, f"{anno['frame_idx']:06d}_{anno['token']}.jpg"
            )

            # go through all bboxes
            bbox_gt_in_ego_list: List[np.ndarray] = list()
            for bbox_index in range(num_obj):
                box_7dof_lidar: np.ndarray = anno["gt_boxes"][bbox_index, :]
                box3d: Box3D = Box3D.array2bbox_xyzlwhyaw(box_7dof_lidar)
                box_in_lidar: np.ndarray = Box3D.box2corners3d_lidar(box3d)  # 8 x 4
                box_in_ego: np.ndarray = box_transform.box_in_lidar_to_ego(
                    box_in_lidar, lidar2ego=anno["lidar2ego"],
                )  # 8 x 4
                bbox_gt_in_ego_list.append(box_in_ego)

            # add ego box
            box3d: Box3D = Box3D.array2bbox_xyzlwhyaw(anno["gt_sdc_bbox_lidar"][:7])
            sdc_box_in_lidar: np.ndarray = Box3D.box2corners3d_lidar(box3d)  # 8 x 4
            sdc_box_in_ego: np.ndarray = box_transform.box_in_lidar_to_ego(
                sdc_box_in_lidar, lidar2ego=anno["lidar2ego"],
            )  # 8 x 4
            bbox_gt_in_ego_list.append(sdc_box_in_ego)

            # aggregate traj
            bbox_fut_lidar_all = np.concatenate(
                (anno["gt_fut_bbox_lidar"], anno["gt_fut_bbox_sdc_lidar"]), axis=0
            )  # N x 12 x 9
            bbox_fut_mask_all = np.concatenate(
                (anno["gt_fut_bbox_mask"], anno["gt_fut_bbox_sdc_mask"]), axis=0
            )  # N x 12 x 1
            bbox_pre_lidar_all = np.concatenate(
                (anno["gt_pre_bbox_lidar"], anno["gt_pre_bbox_sdc_lidar"]), axis=0
            )  # N x 12 x 9
            bbox_pre_mask_all = np.concatenate(
                (anno["gt_pre_bbox_mask"], anno["gt_pre_bbox_sdc_mask"]), axis=0
            )  # N x 12 x 1

            # visualization
            box_vis.vis_box_on_lidar(
                anno["lidar_pts"],
                save_path,
                bbox_gt_in_ego_list=bbox_gt_in_ego_list,
                bev=True,
                id_list=obj_ids + [-1],
                fut_traj=bbox_fut_lidar_all,
                fut_traj_mask=bbox_fut_mask_all,
                pre_traj=bbox_pre_lidar_all,
                pre_traj_mask=bbox_pre_mask_all,
                left_hand=False,
                ego2lidar=np.linalg.inv(anno["lidar2ego"]),
            )

        # delete lidar points in the dicionary to reduce file size prior to saving
        del anno["lidar_pts"]

    return annos


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--data_root",
        type=str,
        default="/mnt/hdd2/datasets/navsim/openscene-v1.1",
        help="root directory of raw carla data",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="/mnt/hdd2/datasets/navsim/openscene-v1.1/infos",
        help="path of output file",
    )
    parser.add_argument(
        "--split",
        type=str,
        default="mini_val",
        help="trainval_train/trainval_val/mini_train/mini_val",
    )
    parser.add_argument(
        "--vis",
        action='store_true',
        help="turn on the visualization",
    )    
    args = parser.parse_args()

    # OpenScenes/nuPlan/NavSim: 
    # mini_train: 43261 (43417 pre-cleaning) -> 6h
    # mini_val: 8440 -> 1.17h
    # val: 115564 (115733 pre-cleaning) -> 16h
    # train: 605263 (607286 pre-cleaning) -> 84h

    # load merged detection data
    data_src = f"{args.data_root}/infos/openscene_{args.split}.pkl"
    with open(data_src, "rb") as f:
        data_src = pickle.load(f)

    seq_index: List[int] = []
    scene_token_pre = None

    for frame_count_global, frame_data in enumerate(data_src):
        # one log can have many scenes/scenarios
        log_name = frame_data["log_name"]
        log_token = frame_data["log_token"]
        scene_name = frame_data["scene_name"]
        scene_token = frame_data["scene_token"]
        token = frame_data["token"]
        frame_idx = frame_data["frame_idx"]

        # print("\nframe_count_global", frame_count_global)
        # print("log_name", log_name)
        # print("log_token", log_token)
        # print("scene_name", scene_name)
        # print("scene_token", scene_token)
        # print("token", token)
        # print("frame_idx", frame_idx)
        # print("frame_idx", frame_idx)
        # print("frame_idx type", type(frame_idx))
        # print(frame_data.keys())
        # DONE: 'token', 'frame_idx', 'timestamp', 'log_name', 'log_token', 'scene_name', 'scene_token',
        # DONE: 'map_location', 'roadblock_ids',
        # 'vehicle_name', 'can_bus',
        # DONE: 'lidar_path',
        # DONE: 'lidar2ego_translation', 'lidar2ego_rotation', 'ego2global_translation', 'ego2global_rotation',
        # 'ego_dynamic_state', 'traffic_lights', 'driving_command', 'cams', 'sample_prev', 'sample_next',
        # DONE: 'ego2global', 'lidar2ego', 'lidar2global',
        # 'anns', 'occ_gt_final_path', 'flow_gt_final_path'
        # zxc

        # if frame_count_global < 9900:
        #     continue

        # process the previous sequence of data, then start a new one now
        if scene_token != scene_token_pre and scene_token_pre is not None:
            # collect temporal data after pre-processing for the entire sequence
            # e.g., add future trajectory into every frame of data
            data_src: List[Dict] = parse_sequence(
                scene_token_pre, data_src, seq_index, args.split, vis=args.vis,
            )
            seq_index: List[int] = []

        # update per-frame data dictionary
        frame_data: Dict = parse_data(frame_data, args.split, vis=args.vis)
        data_src[frame_count_global] = frame_data

        # collect the index of frames that belong to the same sequence
        seq_index.append(frame_count_global)
        scene_token_pre = scene_token

        if frame_count_global % 10 == 0:
            print("finished: ", frame_count_global)

    # process the data clip
    data_src: List[Dict] = parse_sequence(
        scene_token_pre, data_src, seq_index, args.split
    )
    print("total length: ", len(data_src))

    # remove the broken frames
    save_data = []
    for data_dict in data_src:
        if data_dict is not None:
            save_data.append(data_dict)
    print("total length after cleaning: ", len(save_data))

    # save
    data_to_save = {
        "infos": save_data,
        "mapping_tracktoken2globalid": mapping_tracktoken2globalid,
    }
    output_path = os.path.join(args.output_dir, f"nuplan_{args.split}.pkl")
    with open(output_path, "wb") as f:
        pkl.dump(data_to_save, f)