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 .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 Waymo(BaseDataset):
"""This class is used to create a dataset based on the Waymo 3D dataset, and
used in object detection, visualizer, training, or testing.
The Waymo 3D dataset is best suited for autonomous driving applications.
"""
[docs] def __init__(self,
dataset_path,
name='Waymo',
cache_dir='./logs/cache',
use_cache=False,
val_split=3,
**kwargs):
"""Initialize the function by passing the dataset and other details.
Args:
dataset_path: The path to the dataset to use.
name: The name of the dataset (Waymo in this case).
cache_dir: The directory where the cache is stored.
use_cache: Indicates if the dataset should be cached.
val_split: The split value to get a set of images for training, validation, for testing.
Returns:
class: The corresponding class.
"""
super().__init__(dataset_path=dataset_path,
name=name,
cache_dir=cache_dir,
use_cache=use_cache,
val_split=val_split,
**kwargs)
cfg = self.cfg
self.name = cfg.name
self.dataset_path = cfg.dataset_path
self.num_classes = 4
self.label_to_names = self.get_label_to_names()
self.all_files = sorted(
glob(join(cfg.dataset_path, 'velodyne', '*.bin')))
self.train_files = []
self.val_files = []
for f in self.all_files:
idx = Path(f).name.replace('.bin', '')[:3]
idx = int(idx)
if idx < cfg.val_split:
self.train_files.append(f)
else:
self.val_files.append(f)
self.test_files = glob(
join(cfg.dataset_path, 'testing', 'velodyne', '*.bin'))
[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: 'PEDESTRIAN',
1: 'VEHICLE',
2: 'CYCLIST',
3: 'SIGN'
}
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, 6)
[docs] @staticmethod
def read_label(path, calib):
"""Reads labels of bound boxes.
Returns:
The data objects with bound boxes information.
"""
if not Path(path).exists():
return None
with open(path, 'r') as f:
lines = f.readlines()
objects = []
for line in lines:
label = line.strip().split(' ')
center = [float(label[11]), float(label[12]), float(label[13])]
size = [float(label[9]), float(label[8]), float(label[10])]
objects.append(Object3d(center, size, label, calib))
return objects
@staticmethod
def _extend_matrix(mat):
mat = np.concatenate([mat, np.array([[0., 0., 0., 1.]])], axis=0)
return mat
[docs] @staticmethod
def read_calib(path):
"""Reads calibiration for the dataset. You can use them to compare
modeled results to observed results.
Returns:
The camera and the camera image used in calibration.
"""
assert Path(path).exists()
with open(path, 'r') as f:
lines = f.readlines()
obj = lines[0].strip().split(' ')[1:]
P0 = np.array(obj, dtype=np.float32)
obj = lines[1].strip().split(' ')[1:]
P1 = np.array(obj, dtype=np.float32)
obj = lines[2].strip().split(' ')[1:]
P2 = np.array(obj, dtype=np.float32)
obj = lines[3].strip().split(' ')[1:]
P3 = np.array(obj, dtype=np.float32)
obj = lines[4].strip().split(' ')[1:]
P4 = np.array(obj, dtype=np.float32)
obj = lines[5].strip().split(' ')[1:]
R0 = np.array(obj, dtype=np.float32).reshape(3, 3)
rect_4x4 = np.zeros([4, 4], dtype=R0.dtype)
rect_4x4[3, 3] = 1
rect_4x4[:3, :3] = R0
obj = lines[6].strip().split(' ')[1:]
Tr_velo_to_cam = np.array(obj, dtype=np.float32).reshape(3, 4)
Tr_velo_to_cam = Waymo._extend_matrix(Tr_velo_to_cam)
world_cam = np.transpose(rect_4x4 @ Tr_velo_to_cam)
cam_img = np.transpose(P2)
return {'world_cam': world_cam, 'cam_img': cam_img}
[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 WaymoSplit(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'.
"""
cfg = self.cfg
dataset_path = cfg.dataset_path
file_list = []
if split in ['train', 'training']:
# seq_list = cfg.training_split
return self.train_files
elif split in ['test', 'testing']:
return self.test_files
elif split in ['val', 'validation']:
return self.val_files
elif split in ['all']:
return self.train_files + self.val_files + self.test_files
else:
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 WaymoSplit(BaseDatasetSplit):
def __init__(self, dataset, split='train'):
super().__init__(dataset, split=split)
def __len__(self):
return len(self.path_list)
def get_data(self, idx):
pc_path = self.path_list[idx]
label_path = pc_path.replace('velodyne',
'label_all').replace('.bin', '.txt')
calib_path = label_path.replace('label_all', 'calib')
pc = self.dataset.read_lidar(pc_path)
calib = self.dataset.read_calib(calib_path)
label = self.dataset.read_label(label_path, calib)
data = {
'point': pc,
'feat': None,
'calib': calib,
'bounding_boxes': label,
}
return data
def get_attr(self, idx):
pc_path = self.path_list[idx]
name = Path(pc_path).name.split('.')[0]
attr = {'name': name, 'path': pc_path, 'split': self.split}
return attr
class Object3d(BEVBox3D):
"""The class stores details that are object-specific, such as bounding box
coordinates, occlusion and so on.
"""
def __init__(self, center, size, label, calib):
confidence = float(label[15]) if label.__len__() == 16 else -1.0
world_cam = calib['world_cam']
cam_img = calib['cam_img']
# kitti boxes are pointing backwards
yaw = float(label[14]) - np.pi
yaw = yaw - np.floor(yaw / (2 * np.pi) + 0.5) * 2 * np.pi
self.truncation = float(label[1])
self.occlusion = float(
label[2]
) # 0:fully visible 1:partly occluded 2:largely occluded 3:unknown
self.alpha = float(label[3])
self.box2d = np.array((float(label[4]), float(label[5]), float(
label[6]), float(label[7])),
dtype=np.float32)
super().__init__(center, size, yaw, label[0], confidence, world_cam,
cam_img)
self.yaw = float(label[14])
def get_difficulty(self):
"""The method determines difficulty level of the object, such as Easy,
Moderate, or Hard.
"""
height = float(self.box2d[3]) - float(self.box2d[1]) + 1
if height >= 40 and self.truncation <= 0.15 and self.occlusion <= 0:
self.level_str = 'Easy'
return 0 # Easy
elif height >= 25 and self.truncation <= 0.3 and self.occlusion <= 1:
self.level_str = 'Moderate'
return 1 # Moderate
elif height >= 25 and self.truncation <= 0.5 and self.occlusion <= 2:
self.level_str = 'Hard'
return 2 # Hard
else:
self.level_str = 'UnKnown'
return -1
def to_str(self):
print_str = '%s %.3f %.3f %.3f box2d: %s hwl: [%.3f %.3f %.3f] pos: %s ry: %.3f' \
% (self.label_class, self.truncation, self.occlusion, self.alpha, self.box2d, self.size[2],
self.size[0], self.size[1],
self.center, self.yaw)
return print_str
def to_kitti_format(self):
"""This method transforms the class to kitti format."""
kitti_str = '%s %.2f %d %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f' \
% (self.label_class, self.truncation, int(self.occlusion), self.alpha, self.box2d[0], self.box2d[1],
self.box2d[2], self.box2d[3], self.size[2], self.size[0], self.size[1], self.center[0],
self.center[1], self.center[2],
self.yaw)
return kitti_str
DATASET._register_module(Waymo)