490 lines
19 KiB
Python
490 lines
19 KiB
Python
"""
|
|
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'
|
|
)
|
|
|
|
|
|
|
|
|