import numpy as np
import os, argparse, pickle, sys
from os.path import exists, join, isfile, dirname, abspath, split
from pathlib import Path
from glob import glob
import logging
import yaml
from scipy.spatial.transform import Rotation as R
from .base_dataset import BaseDataset, BaseDatasetSplit
from ..utils import Config, make_dir, DATASET
from .utils import BEVBox3D
logging.basicConfig(
level=logging.INFO,
format='%(levelname)s - %(asctime)s - %(module)s - %(message)s',
)
log = logging.getLogger(__name__)
[docs]class NuScenes(BaseDataset):
"""This class is used to create a dataset based on the NuScenes 3D dataset,
and used in object detection, visualizer, training, or testing.
The NuScenes 3D dataset is best suited for autonomous driving applications.
"""
[docs] def __init__(self,
dataset_path,
info_path=None,
name='NuScenes',
cache_dir='./logs/cache',
use_cache=False,
**kwargs):
"""
Initialize the function by passing the dataset and other details.
Args:
dataset_path: The path to the dataset to use.
info_path: The path to the file that includes information about the dataset. This is default to dataset path if nothing is provided.
name: The name of the dataset (NuScenes in this case).
cache_dir: The directory where the cache is stored.
use_cache: Indicates if the dataset should be cached.
Returns:
class: The corresponding class.
"""
if info_path is None:
info_path = dataset_path
super().__init__(dataset_path=dataset_path,
info_path=info_path,
name=name,
cache_dir=cache_dir,
use_cache=use_cache,
**kwargs)
cfg = self.cfg
self.name = cfg.name
self.dataset_path = cfg.dataset_path
self.num_classes = 10
self.label_to_names = self.get_label_to_names()
self.train_info = {}
self.test_info = {}
self.val_info = {}
if os.path.exists(join(info_path, 'infos_train.pkl')):
self.train_info = pickle.load(
open(join(info_path, 'infos_train.pkl'), 'rb'))
if os.path.exists(join(info_path, 'infos_val.pkl')):
self.val_info = pickle.load(
open(join(info_path, 'infos_val.pkl'), 'rb'))
if os.path.exists(join(info_path, 'infos_test.pkl')):
self.test_info = pickle.load(
open(join(info_path, 'infos_test.pkl'), 'rb'))
[docs] @staticmethod
def get_label_to_names():
"""
Returns a label to names dictonary object.
Returns:
A dict where keys are label numbers and
values are the corresponding names.
"""
label_to_names = {
0: 'ignore',
1: 'barrier',
2: 'bicycle',
3: 'bus',
4: 'car',
5: 'construction_vehicle',
6: 'motorcycle',
7: 'pedestrian',
8: 'traffic_cone',
9: 'trailer',
10: 'truck'
}
return label_to_names
[docs] @staticmethod
def read_lidar(path):
"""
Reads lidar data from the path provided.
Returns:
A data object with lidar information.
"""
assert Path(path).exists()
return np.fromfile(path, dtype=np.float32).reshape(-1, 5)
[docs] @staticmethod
def read_label(info, calib):
"""
Reads labels of bound boxes.
Returns:
The data objects with bound boxes information.
"""
mask = info['num_lidar_pts'] != 0
boxes = info['gt_boxes'][mask]
names = info['gt_names'][mask]
objects = []
for name, box in zip(names, boxes):
center = [float(box[0]), float(box[1]), float(box[2])]
size = [float(box[3]), float(box[5]), float(box[4])]
ry = float(box[6])
yaw = ry - np.pi
yaw = yaw - np.floor(yaw / (2 * np.pi) + 0.5) * 2 * np.pi
world_cam = calib['world_cam']
objects.append(BEVBox3D(center, size, yaw, name, -1.0, world_cam))
objects[-1].yaw = ry
return objects
[docs] def get_split(self, split):
"""Returns a dataset split.
Args:
split: A string identifying the dataset split that is usually one of
'training', 'test', 'validation', or 'all'.
Returns:
A dataset split object providing the requested subset of the data.
"""
return NuSceneSplit(self, split=split)
[docs] def get_split_list(self, split):
"""Returns the list of data splits available.
Args:
split: A string identifying the dataset split that is usually one of
'training', 'test', 'validation', or 'all'.
Returns:
A dataset split object providing the requested subset of the data.
Raises:
ValueError: Indicates that the split name passed is incorrect. The split name should be one of
'training', 'test', 'validation', or 'all'.
"""
if split in ['train', 'training']:
return self.train_info
elif split in ['test', 'testing']:
return self.test_info
elif split in ['val', 'validation']:
return self.val_info
raise ValueError("Invalid split {}".format(split))
[docs] def is_tested(self, attr):
"""Checks if a datum in the dataset has been tested.
Args:
dataset: The current dataset to which the datum belongs to.
attr: The attribute that needs to be checked.
Returns:
If the dataum attribute is tested, then resturn the path where the attribute is stored; else, returns false.
"""
pass
[docs] def save_test_result(self, results, attr):
"""Saves the output of a model.
Args:
results: The output of a model for the datum associated with the attribute passed.
attr: The attributes that correspond to the outputs passed in results.
"""
pass
class NuSceneSplit(BaseDatasetSplit):
def __init__(self, dataset, split='train'):
super().__init__(dataset, split=split)
self.infos = self.path_list
self.path_list = []
for info in self.infos:
self.path_list.append(info['lidar_path'])
def __len__(self):
return len(self.infos)
def get_data(self, idx):
info = self.infos[idx]
lidar_path = info['lidar_path']
world_cam = np.eye(4)
world_cam[:3, :3] = R.from_quat(info['lidar2ego_rot']).as_matrix()
world_cam[:3, -1] = info['lidar2ego_tr']
calib = {'world_cam': world_cam.T}
pc = self.dataset.read_lidar(lidar_path)
label = self.dataset.read_label(info, calib)
data = {
'point': pc,
'feat': None,
'calib': calib,
'bounding_boxes': label,
}
return data
def get_attr(self, idx):
info = self.infos[idx]
pc_path = info['lidar_path']
name = Path(pc_path).name.split('.')[0]
attr = {'name': name, 'path': str(pc_path), 'split': self.split}
return attr
DATASET._register_module(NuScenes)