pandaset_dataset.py

"""
    Dataset from Pandaset (Hesai)
"""

import pickle
import os
try:
    import pandas as pd
    import pandaset as ps
except:
    pass 
import numpy as np

from ..dataset import DatasetTemplate
from ...ops.roiaware_pool3d import roiaware_pool3d_utils

import torch


def pose_dict_to_numpy(pose):
    """
        Conert pandaset pose dict to a numpy vector in order to pass it through the network
    """
    pose_np = [pose["position"]["x"],
               pose["position"]["y"],
               pose["position"]["z"],
               pose["heading"]["w"],
               pose["heading"]["x"],
               pose["heading"]["y"],
               pose["heading"]["z"]]

    return pose_np


def pose_numpy_to_dict(pose):
    """
        Conert pandaset pose dict to a numpy vector in order to pass it through the network
    """
    pose_dict = {'position':
                    {'x': pose[0],
                     'y': pose[1],
                     'z': pose[2]},
                 'heading':
                    {'w': pose[3],
                     'x': pose[4],
                     'y': pose[5],
                     'z': pose[6]}}

    return pose_dict


class PandasetDataset(DatasetTemplate):
    def __init__(self, dataset_cfg, class_names, training=True, root_path=None, logger=None):
        """
        Args:
            root_path:
            dataset_cfg:
            class_names:
            training:
            logger:
        """
        super().__init__(
            dataset_cfg=dataset_cfg, class_names=class_names, training=training, root_path=root_path, logger=logger
        )
        if root_path is None:
            root_path = self.dataset_cfg.DATA_PATH
        self.dataset = ps.DataSet(os.path.join(root_path, 'dataset'))
        self.split = self.dataset_cfg.DATA_SPLIT[self.mode]
        self.pandaset_infos = []
        self.include_pandaset_infos(self.mode)


    def include_pandaset_infos(self, mode):
        if self.logger is not None:
            self.logger.info('Loading PandaSet dataset')
        pandaset_infos = []

        for info_path in self.dataset_cfg.INFO_PATH[mode]:
            info_path = os.path.join(self.root_path, info_path)
            if not os.path.exists(info_path):
                continue
            with open(info_path, 'rb') as f:
                infos = pickle.load(f)
                pandaset_infos.extend(infos)

        self.pandaset_infos.extend(pandaset_infos)

        if self.logger is not None:
            self.logger.info('Total samples for PandaSet dataset ({}): {}'.format(self.mode, len(pandaset_infos)))


    def set_split(self, split):
        self.sequences = self.dataset_cfg.SEQUENCES[split]
        self.split = split


    def __len__(self):
        return len(self.pandaset_infos)


    def __getitem__(self, index):
        """
        To support a custom dataset, implement this function to load the raw data (and labels), then transform them to
        the unified normative coordinate (x pointing forward, z pointing upwards) and call the function self.prepare_data() to process the data and send them
        to the model.

        Args:
            index:

        Returns:

        """
        info = self.pandaset_infos[index]
        seq_idx = info['sequence']

        pose = self._get_pose(info)
        points = self._get_lidar_points(info, pose)
        boxes, labels, zrot_world_to_ego = self._get_annotations(info, pose)
        pose_np = pose_dict_to_numpy(pose)

        input_dict = {'points': points,
                      'gt_boxes': boxes,
                      'gt_names': labels,
                      'sequence': int(seq_idx),
                      'frame_idx': info['frame_idx'],
                      'zrot_world_to_ego': zrot_world_to_ego,
                      'pose': pose_dict_to_numpy(pose)
                     }
        # seq_idx is converted to int because strings can't be passed to
        # the gpu in pytorch
        # zrot_world_to_ego is propagated in order to be able to transform the
        # predicted yaws back to world coordinates

        data_dict = self.prepare_data(data_dict=input_dict)

        return data_dict


    def _get_pose(self, info):
        seq_idx = info['sequence']
        # get pose for world to ego frame transformation
        if self.dataset[seq_idx].lidar.poses is None:
            self.dataset[seq_idx].lidar._load_poses()

        pose = self.dataset[seq_idx].lidar.poses[info['frame_idx']]

        return pose


    def _get_lidar_points(self, info, pose):
        """
        Get lidar in the unified normative coordinate system for a given frame
        The intensity is normalized to fit [0-1] range (pandaset intensity is in [0-255] range)
        """
        # get lidar points
        lidar_frame = pd.read_pickle(info['lidar_path'])
        # get points for the required lidar(s) only
        device = self.dataset_cfg.get('LIDAR_DEVICE', 0)
        if device != -1:
            lidar_frame = lidar_frame[lidar_frame.d == device]
        world_points = lidar_frame.to_numpy()
        # There seems to be issues with the automatic deletion of pandas datasets sometimes
        del lidar_frame

        points_loc = world_points[:, :3]
        points_int = world_points[:, 3]

        # nromalize intensity
        points_int = points_int / 255

        ego_points = ps.geometry.lidar_points_to_ego(points_loc, pose)
        # Pandaset ego coordinates are:
        # - x pointing to the right
        # - y pointing to the front
        # - z pointing up
        # Normative coordinates are:
        # - x pointing foreward
        # - y pointings to the left
        # - z pointing to the top
        # So a transformation is required to the match the normative coordinates
        ego_points = ego_points[:, [1, 0, 2]] # switch x and y
        ego_points[:, 1] = - ego_points[:, 1] # revert y axis

        return np.append(ego_points, np.expand_dims(points_int, axis=1), axis=1).astype(np.float32)


    def _get_annotations(self,info, pose):
        """
        Get box informations in the unified normative coordinate system for a given frame
        """

        # get boxes
        cuboids = pd.read_pickle(info["cuboids_path"])
        device = self.dataset_cfg.get('LIDAR_DEVICE', 0)
        if device != -1:
            # keep cuboids that are seen by a given device
            cuboids = cuboids[cuboids["cuboids.sensor_id"] != 1 - device]

        xs = cuboids['position.x'].to_numpy()
        ys = cuboids['position.y'].to_numpy()
        zs = cuboids['position.z'].to_numpy()
        dxs = cuboids['dimensions.x'].to_numpy()
        dys = cuboids['dimensions.y'].to_numpy()
        dzs = cuboids['dimensions.z'].to_numpy()
        yaws = cuboids['yaw'].to_numpy()
        labels = cuboids['label'].to_numpy()

        del cuboids  # There seem to be issues with the automatic deletion of pandas datasets sometimes

        labels = np.array([self.dataset_cfg.TRAINING_CATEGORIES.get(lab, lab)
                           for lab in labels] )

        # Compute the center points coordinates in ego coordinates
        centers = np.vstack([xs, ys, zs]).T
        ego_centers = ps.geometry.lidar_points_to_ego(centers, pose)

        # Compute the yaw in ego coordinates
        # The following implementation supposes that the pitch of the car is
        # negligible compared to its yaw, in order to be able to express the
        # bbox coordinates in the ego coordinate system with an {axis aligned
        # box + yaw} only representation
        yaxis_points_from_pose = ps.geometry.lidar_points_to_ego(np.array([[0, 0, 0], [0, 1., 0]]), pose)
        yaxis_from_pose = yaxis_points_from_pose[1, :] - yaxis_points_from_pose[0, :]

        if yaxis_from_pose[-1] >= 10**-1:
            if self.logger is not None:
                self.logger.warning("The car's pitch is supposed to be negligible " +
                                    "sin(pitch) is >= 10**-1 ({})".format(yaxis_from_pose[-1]))

        # rotation angle in rads of the y axis around thz z axis
        zrot_world_to_ego = np.arctan2(-yaxis_from_pose[0], yaxis_from_pose[1])
        ego_yaws = yaws + zrot_world_to_ego

        # Pandaset ego coordinates are:
        # - x pointing to the right
        # - y pointing to the front
        # - z pointing up
        # Normative coordinates are:
        # - x pointing foreward
        # - y pointings to the left
        # - z pointing to the top
        # So a transformation is required to the match the normative coordinates
        ego_xs = ego_centers[:, 1]
        ego_ys = -ego_centers[:, 0]
        ego_zs = ego_centers[:, 2]
        ego_dxs = dys
        ego_dys = dxs  # stays >= 0
        ego_dzs = dzs

        ego_boxes = np.vstack([ego_xs, ego_ys, ego_zs, ego_dxs, ego_dys, ego_dzs, ego_yaws]).T

        return ego_boxes.astype(np.float32), labels, zrot_world_to_ego


    @staticmethod
    def generate_prediction_dicts(batch_dict, pred_dicts, class_names, output_path=None):
        """
        To support a custom dataset, implement this function to receive the predicted results from the model, and then
        transform the unified normative coordinate to your required coordinate, and optionally save them to disk.

        Args:
            batch_dict: dict of original data from the dataloader
            pred_dicts: dict of predicted results from the model
                pred_boxes: (N, 7), Tensor
                pred_scores: (N), Tensor
                pred_labels: (N), Tensor
            class_names:
            output_path: if it is not None, save the results to this path
        Returns:

        """

        def generate_single_sample_dataframe(batch_index, box_dict, zrot_world_to_ego, pose):
            pred_boxes = box_dict["pred_boxes"].cpu().numpy()
            pred_scores = box_dict["pred_scores"].cpu().numpy()
            pred_labels = box_dict["pred_labels"].cpu().numpy()
            zrot = zrot_world_to_ego.cpu().numpy()
            pose_dict = pose_numpy_to_dict(pose.cpu().numpy())

            xs = pred_boxes[:, 0]
            ys = pred_boxes[:, 1]
            zs = pred_boxes[:, 2]
            dxs = pred_boxes[:, 3]
            dys = pred_boxes[:, 4]
            dzs = pred_boxes[:, 5]
            yaws = pred_boxes[:, 6]
            names = np.array(class_names)[pred_labels - 1]  # Predicted labels start on 1

            # convert from normative coordinates to pandaset ego coordinates
            ego_xs = - ys
            ego_ys = xs
            ego_zs = zs
            ego_dxs = dys
            ego_dys = dxs
            ego_dzs = dzs
            ego_yaws = yaws

            # convert from pandaset ego coordinates to world coordinates
            # for the moment, an simplified estimation of the ego yaw is computed in __getitem__
            # which sets ego_yaw = world_yaw + zrot_world_to_ego
            world_yaws = ego_yaws - zrot

            ego_centers = np.vstack([ego_xs, ego_ys, ego_zs]).T
            world_centers = ps.geometry.ego_to_lidar_points(ego_centers, pose_dict)
            world_xs = world_centers[:, 0]
            world_ys = world_centers[:, 1]
            world_zs = world_centers[:, 2]
            # dx, dy, dz remain unchanged as the bbox orientation is handled by
            # the yaw information

            data_dict = {'position.x': world_xs,
                         'position.y': world_ys,
                         'position.z': world_zs,
                         'dimensions.x': ego_dxs,
                         'dimensions.y': ego_dys,
                         'dimensions.z': ego_dzs,
                         'yaw': world_yaws % (2 * np.pi),
                         'label': names,
                         'score': pred_scores
            }

            return pd.DataFrame(data_dict)


        annos = []
        for index, box_dict in enumerate(pred_dicts):
            frame_idx = batch_dict['frame_idx'][index]
            seq_idx = batch_dict['sequence'][index]
            zrot = batch_dict['zrot_world_to_ego'][index]
            pose = batch_dict['pose'][index]

            single_pred_df = generate_single_sample_dataframe(index, box_dict, zrot, pose)


            single_pred_dict = {'preds' : single_pred_df,
                                # 'name 'ensures testing the number of detections in a compatible format as kitti
                                'name' : single_pred_df['label'].tolist(),
                                'frame_idx': frame_idx,
                                'sequence': str(seq_idx).zfill(3)}
            # seq_idx was converted to int in self.__getitem__` because strings
            # can't be passed to the gpu in pytorch.
            # To convert it back to a string, we assume that the sequence is
            # provided in pandaset format with 3 digits

            if output_path is not None:
                frame_id = str(int(frame_idx)).zfill(2)
                seq_id = str(int(seq_idx)).zfill(3)
                cur_det_file = os.path.join(output_path, seq_id, 'predictions',
                                            'cuboids', ("{}.pkl.gz".format(frame_id)))
                os.makedirs(os.path.dirname(cur_det_file), exist_ok=True)
                single_pred_df.to_pickle(cur_det_file)

            annos.append(single_pred_dict)

        return annos


    def get_infos(self):
        """
        Generate the dataset infos dict for each sample of the dataset.
        For each sample, this dict contains:
            - the sequence index
            - the frame index
            - the path to the lidar data
            - the path to the bounding box annotations
        """
        infos = []
        for seq in self.sequences:
            s = self.dataset[seq]
            s.load_lidar()
            if len(s.lidar.data) > 100:
                raise ValueError("The implementation for this dataset assumes that each sequence is " +
                                 "no longer than 100 frames. The current sequence has {}".format(len(s.lidar.data)))
            info = [{'sequence': seq,
                     'frame_idx': ii,
                     'lidar_path': os.path.join(self.root_path, 'dataset', seq, 'lidar', ("{:02d}.pkl.gz".format(ii))),
                     'cuboids_path': os.path.join(self.root_path, 'dataset', seq,
                                                  'annotations', 'cuboids', ("{:02d}.pkl.gz".format(ii)))
                    } for ii in range(len(s.lidar.data))]
            infos.extend(info)
            del self.dataset._sequences[seq]

        return infos


    def create_groundtruth_database(self, info_path=None, used_classes=None, split='train'):
        database_save_path = os.path.join(self.root_path,
                'gt_database' if split == 'train' else 'gt_database_{}'.format(split))
        db_info_save_path = os.path.join(self.root_path,
                'pandaset_dbinfos_{}.pkl'.format(split))

        os.makedirs(database_save_path, exist_ok=True)
        all_db_infos = {}

        with open(info_path, 'rb') as f:
            infos = pickle.load(f)

        for k in range(len(infos)):
            print('gt_database sample: %d/%d' % (k + 1, len(infos)))
            info = infos[k]
            sample_idx = info['frame_idx']
            pose = self._get_pose(info)
            points = self._get_lidar_points(info, pose)
            gt_boxes, names, _ = self._get_annotations(info, pose)

            num_obj = gt_boxes.shape[0]

            point_indices = roiaware_pool3d_utils.points_in_boxes_cpu(
                torch.from_numpy(points[:, 0:3]), torch.from_numpy(gt_boxes)
            ).numpy()  # (nboxes, npoints)

            for i in range(num_obj):
                tmp_name = names[i].replace("/", "").replace(" ", "")
                filename = '%s_%s_%d.bin' % (sample_idx, tmp_name, i)
                filepath = os.path.join(database_save_path, filename)
                gt_points = points[point_indices[i] > 0]
                gt_points[:, :3] -= gt_boxes[i, :3]
                with open(filepath, 'wb') as f:
                    gt_points.tofile(f)

                if (used_classes is None) or names[i] in used_classes:
                    db_path = os.path.relpath(filepath, self.root_path)  # gt_database/xxxxx.bin
                    db_info = {'name': names[i], 'path': db_path, 'gt_idx': i,
                               'box3d_lidar': gt_boxes[i], 'num_points_in_gt': gt_points.shape[0],
                               'difficulty': -1}
                    if names[i] in all_db_infos:
                        all_db_infos[names[i]].append(db_info)
                    else:
                        all_db_infos[names[i]] = [db_info]
        for k, v in all_db_infos.items():
            print('Database %s: %d' % (k, len(v)))

        with open(db_info_save_path, 'wb') as f:
            pickle.dump(all_db_infos, f)


    def evaluation(self, det_annos, class_names, **kwargs):
        self.logger.warning('Evaluation is not implemented for Pandaset as there is no official one. ' +
                            'Returning an empty evaluation result.')
        ap_result_str = ''
        ap_dict = {}

        return ap_result_str, ap_dict


def create_pandaset_infos(dataset_cfg, class_names, data_path, save_path):
    """
    Create dataset_infos files in order not to have it in a preprocessed pickle
    file with the info for each sample
    See PandasetDataset.get_infos for further details.
    """
    dataset = PandasetDataset(dataset_cfg=dataset_cfg, class_names=class_names, root_path=data_path, training=False)
    for split in ["train", "val", "test"]:
        print("---------------- Start to generate {} data infos ---------------".format(split))
        dataset.set_split(split)
        infos = dataset.get_infos()
        file_path = os.path.join(save_path, 'pandaset_infos_{}.pkl'.format(split))
        with open(file_path, 'wb') as f:
            pickle.dump(infos, f)
        print("Pandaset info {} file is saved to {}".format(split, file_path))

    print('------------Start create groundtruth database for data augmentation-----------')
    dataset = PandasetDataset(dataset_cfg=dataset_cfg, class_names=class_names, root_path=data_path, training=False)
    dataset.set_split("train")
    dataset.create_groundtruth_database(
        os.path.join(save_path, 'pandaset_infos_train.pkl'),
        split="train"
    )
    print('---------------Data preparation Done---------------')


if __name__ == '__main__':
    import sys
    if sys.argv.__len__() > 1 and sys.argv[1] == 'create_pandaset_infos':
        import yaml
        from pathlib import Path
        from easydict import EasyDict
        dataset_cfg = EasyDict(yaml.safe_load(open(sys.argv[2])))
        ROOT_DIR = (Path(__file__).resolve().parent / '../../../').resolve()
        create_pandaset_infos(
            dataset_cfg=dataset_cfg,
            class_names=['Car', 'Pedestrian', 'Cyclist'],
            data_path=ROOT_DIR / 'data' / 'pandaset',
            save_path=ROOT_DIR / 'data' / 'pandaset'
        )