from typing import Tuple, List
from cv2 import repeat
from einops import rearrange, repeat
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import torch
from torch import Tensor
import numpy as np
import pickle
import os
from mmcv import Config
class WaymoDataset(Dataset):
"""
Dataset loader for custom preprocessed files of Waymo data.
Args:
path: path to the dataset directory
args: global settings
"""
def __init__(self, cfg: Config, split: str, input_angle: bool = True):
super(WaymoDataset, self).__init__()
self.p_exchange_two_first = 1
if "val" in split.lower():
path = cfg.val_dataset_path
elif "test" in split.lower():
path = cfg.test_dataset_path
elif "sample" in split.lower():
path = cfg.sample_dataset_path
else:
path = cfg.train_dataset_path
self.p_exchange_two_first = cfg.p_exchange_two_first
self.file_list = [
os.path.join(path, name)
for name in os.listdir(path)
if os.path.isfile(os.path.join(path, name))
]
self.normalize = cfg.normalize_angle
# self.load_dataset(path, 16)
# self.idx_list = list(self.dataset.keys())
self.input_angle = input_angle
self.hist_len = cfg.num_steps
self.fut_len = cfg.num_steps_future
self.time_len = self.hist_len + self.fut_len
self.min_num_obs = cfg.min_num_observation
self.max_size_lane = cfg.max_size_lane
self.random_rotation = cfg.random_rotation
self.random_translation = cfg.random_translation
self.angle_std = cfg.angle_std
self.translation_distance_std = cfg.translation_distance_std
self.max_num_agents = cfg.max_num_agents
self.max_num_objects = cfg.max_num_objects
self.state_dim = cfg.state_dim
self.map_state_dim = cfg.map_state_dim
self.dt = cfg.dt
if "val" in os.path.basename(path).lower():
self.dataset_size_limit = cfg.val_dataset_size_limit
else:
self.dataset_size_limit = cfg.train_dataset_size_limit
def __len__(self):
if self.dataset_size_limit is not None:
return min(len(self.file_list), self.dataset_size_limit)
else:
return len(self.file_list)
def __getitem__(
self, idx: int
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""
Get the item at index idx in the dataset. Normalize the scene and output absolute angle and position.
Returns:
trajectories, mask, mask_loss, lanes, mask_lanes, angle, mean_position
"""
selected_file = self.file_list[idx]
with open(selected_file, "rb") as handle:
dataset = pickle.load(handle)
rel_state_all = dataset["traj"]
mask_all = dataset["mask_traj"]
mask_loss = dataset["mask_to_predict"]
rel_lane_all = dataset["lanes"]
mask_lane_all = dataset["mask_lanes"]
mean_pos = dataset["mean_pos"]
assert (
(
rel_state_all[self.hist_len + 5 :, :, :2][mask_all[self.hist_len + 5 :]]
!= 0
)
.any(-1)
.all()
)
assert (
(
rel_state_all[self.hist_len + 5 :, :, :2][
mask_loss[self.hist_len + 5 :]
]
!= 0
)
.any(-1)
.all()
)
if "lane_states" in dataset.keys():
lane_states = dataset["lane_states"]
else:
lane_states = None
if np.random.rand() > self.p_exchange_two_first:
rel_state_all[:, [0, 1]] = rel_state_all[:, [1, 0]]
mask_all[:, [0, 1]] = mask_all[:, [1, 0]]
mask_loss[:, [0, 1]] = mask_loss[:, [1, 0]]
assert (
(
rel_state_all[self.hist_len + 5 :, :, :2][mask_all[self.hist_len + 5 :]]
!= 0
)
.any(-1)
.all()
)
assert (
(
rel_state_all[self.hist_len + 5 :, :, :2][
mask_loss[self.hist_len + 5 :]
]
!= 0
)
.any(-1)
.all()
)
if self.normalize:
angle = rel_state_all[self.hist_len - 1, 1, 2]
if self.random_rotation:
if self.normalize:
angle += np.random.normal(0, self.angle_std)
else:
angle += np.random.uniform(-np.pi, np.pi)
if self.random_translation:
distance = (
np.random.normal([0, 0], self.translation_distance_std, 2)
* mask_all[self.hist_len - 1 : self.hist_len, :, None]
- rel_state_all[self.hist_len - 1 : self.hist_len, 1:2, :2]
)
else:
distance = -rel_state_all[self.hist_len - 1 : self.hist_len, 1:2, :2]
rel_state_all[:, :, :2] += distance
rel_lane_all[:, :, :2] += distance
mean_pos += distance[0, 0, :]
rel_state_all = self.scene_rotation(rel_state_all, -angle)
rel_lane_all = self.scene_rotation(rel_lane_all, -angle)
else:
if self.random_translation:
distance = np.random.normal([0, 0], self.translation_distance_std, 2)
rel_state_all = (
rel_state_all
+ mask_all[self.hist_len - 1 : self.hist_len, :, None] * distance
)
rel_lane_all = (
rel_lane_all
+ mask_all[self.hist_len - 1 : self.hist_len, :, None] * distance
)
if self.random_rotation:
angle = np.random.uniform(0, 2 * np.pi)
rel_state_all = self.scene_rotation(rel_state_all, angle)
rel_lane_all = self.scene_rotation(rel_lane_all, angle)
else:
angle = 0
return (
rel_state_all,
mask_all,
mask_loss,
rel_lane_all,
mask_lane_all,
lane_states,
angle,
mean_pos,
idx,
)
@staticmethod
def scene_rotation(coor: np.ndarray, angle: float) -> np.ndarray:
"""
Rotate all the coordinates with the same angle
Args:
coor: array of x, y coordinates
angle: radiants to rotate the coordinates by
Returns:
coor_rotated
"""
rot_matrix = np.zeros((2, 2))
c = np.cos(angle)
s = np.sin(angle)
rot_matrix[0, 0] = c
rot_matrix[0, 1] = -s
rot_matrix[1, 0] = s
rot_matrix[1, 1] = c
coor[..., :2] = np.matmul(
rot_matrix, np.expand_dims(coor[..., :2], axis=-1)
).squeeze(-1)
if coor.shape[-1] > 2:
coor[..., 2] += angle
if coor.shape[-1] >= 5:
coor[..., 3:5] = np.matmul(
rot_matrix, np.expand_dims(coor[..., 3:5], axis=-1)
).squeeze(-1)
return coor
def fill_past(self, past, mask_past):
current_velocity = past[..., 0, 3:5]
for t in range(1, past.shape[-2]):
current_velocity = torch.where(
mask_past[..., t, None], past[..., t, 3:5], current_velocity
)
past[..., t, 3:5] = current_velocity
predicted_position = past[..., t - 1, :2] + current_velocity * self.dt
past[..., t, :2] = torch.where(
mask_past[..., t, None], past[..., t, :2], predicted_position
)
return past
def collate_fn(
self, samples: List
) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
"""
Assemble trajectories into batches with 0-padding.
Args:
samples: list of sampled trajectories (list of outputs of __getitem__)
Returns:
(starred dimensions have different values from one batch to the next but the ones with the same name are consistent within the batch)
batch : ((batch_size, num_agents*, num_steps, state_dim), # past trajectories of all agents in the scene
(batch_size, num_agents*, num_steps), # mask past False where past trajectories are padding data
(batch_size, num_agents*, num_steps_future, state_dim), # future trajectory
(batch_size, num_agents*, num_steps_future), # mask future False where future trajectories are padding data
(batch_size, num_agents*, num_steps_future), # mask loss False where future trajectories are not to be predicted
(batch_size, num_objects*, object_seq_len*, map_state_dim),# map object sequences in the scene
(batch_size, num_objects*, object_seq_len*), # mask map False where map objects are padding data
(batch_size, num_agents*, state_dim), # position offset of all agents relative to ego at present time
(batch_size, num_steps, state_dim), # ego past trajectory
(batch_size, num_steps_future, state_dim)) # ego future trajectory
"""
max_n_vehicle = 50
max_n_lanes = 0
for (
coor,
mask,
mask_loss,
lanes,
mask_lanes,
lane_states,
mean_angle,
mean_pos,
idx,
) in samples:
# time_len_coor = self._count_last_obs(coor, hist_len)
# num_vehicle = np.sum(time_len_coor > self.min_num_obs)
num_vehicle = coor.shape[1]
num_lanes = lanes.shape[1]
max_n_vehicle = max(num_vehicle, max_n_vehicle)
max_n_lanes = max(num_lanes, max_n_lanes)
if max_n_vehicle <= 0:
raise RuntimeError
data_batch = np.zeros(
[self.time_len, len(samples), max_n_vehicle, self.state_dim]
)
mask_batch = np.zeros([self.time_len, len(samples), max_n_vehicle])
mask_loss_batch = np.zeros([self.time_len, len(samples), max_n_vehicle])
lane_batch = np.zeros(
[self.max_size_lane, len(samples), max_n_lanes, self.map_state_dim]
)
mask_lane_batch = np.zeros([self.max_size_lane, len(samples), max_n_lanes])
mean_angle_batch = np.zeros([len(samples)])
mean_pos_batch = np.zeros([len(samples), 2])
tag_list = np.zeros([len(samples)])
idx_list = [0 for _ in range(len(samples))]
for sample_ind, (
coor,
mask,
mask_loss,
lanes,
mask_lanes,
lane_states,
mean_angle,
mean_pos,
idx,
) in enumerate(samples):
data_batch[:, sample_ind, : coor.shape[1], :] = coor[: self.time_len, :, :]
mask_batch[:, sample_ind, : mask.shape[1]] = mask[: self.time_len, :]
mask_loss_batch[:, sample_ind, : mask.shape[1]] = mask_loss[
: self.time_len, :
]
lane_batch[: lanes.shape[0], sample_ind, : lanes.shape[1], :2] = lanes
if lane_states is not None:
lane_states = repeat(
lane_states[:, : self.hist_len],
"objects time features -> one objects (time features)",
one=1,
)
lane_batch[
: lanes.shape[0], sample_ind, : lanes.shape[1], 2:
] = lane_states
mask_lane_batch[
: mask_lanes.shape[0], sample_ind, : mask_lanes.shape[1]
] = mask_lanes
mean_angle_batch[sample_ind] = mean_angle
mean_pos_batch[sample_ind, :] = mean_pos
# tag_list[sample_ind] = self.dataset[idx]["tag"]
idx_list[sample_ind] = idx
data_batch = torch.from_numpy(data_batch.astype("float32"))
mask_batch = torch.from_numpy(mask_batch.astype("bool"))
lane_batch = torch.from_numpy(lane_batch.astype("float32"))
mask_lane_batch = torch.from_numpy(mask_lane_batch.astype("bool"))
mean_pos_batch = torch.from_numpy(mean_pos_batch.astype("float32"))
mask_loss_batch = torch.from_numpy(mask_loss_batch.astype("bool"))
data_batch = rearrange(
data_batch, "time batch agents features -> batch agents time features"
)
mask_batch = rearrange(mask_batch, "time batch agents -> batch agents time")
mask_loss_batch = rearrange(
mask_loss_batch, "time batch agents -> batch agents time"
)
lane_batch = rearrange(
lane_batch,
"object_seq_len batch objects features-> batch objects object_seq_len features",
)
mask_lane_batch = rearrange(
mask_lane_batch,
"object_seq_len batch objects -> batch objects object_seq_len",
)
# The two first agents are the ones interacting, others are sorted by distance from the first agent
# Objects are also sorted by distance from the first agent
# Therefore, the limits in number, max_num_agents and max_num_objects can be seen as adaptative distance limits.
if not self.input_angle:
data_batch = torch.cat((data_batch[..., :2], data_batch[..., 3:]), dim=-1)
traj_past = data_batch[:, : self.max_num_agents, : self.hist_len, :]
mask_past = mask_batch[:, : self.max_num_agents, : self.hist_len]
traj_fut = data_batch[
:, : self.max_num_agents, self.hist_len : self.hist_len + self.fut_len, :
]
mask_fut = mask_batch[
:, : self.max_num_agents, self.hist_len : self.hist_len + self.fut_len
]
ego_past = data_batch[:, 0:1, : self.hist_len, :]
ego_fut = data_batch[:, 0:1, self.hist_len : self.hist_len + self.fut_len, :]
lane_batch = lane_batch[:, : self.max_num_objects]
mask_lane_batch = mask_lane_batch[:, : self.max_num_objects]
# Define what to predict (could be from Waymo's label of what to predict or the other agent that interacts with the ego...)
mask_loss_batch = torch.logical_and(
mask_loss_batch[
:, : self.max_num_agents, self.hist_len : self.hist_len + self.fut_len
],
mask_past.any(-1, keepdim=True),
)
# Remove all other agents so the model should only predict the first one
mask_loss_batch[:, 0] = False
mask_loss_batch[:, 2:] = False
# Normalize...
# traj_past = self.fill_past(traj_past, mask_past)
dynamic_state_size = 5 if self.input_angle else 4
offset_batch = traj_past[..., -1, :dynamic_state_size].clone()
traj_past[..., :dynamic_state_size] = traj_past[
..., :dynamic_state_size
] - offset_batch.unsqueeze(-2)
traj_fut[..., :dynamic_state_size] = traj_fut[
..., :dynamic_state_size
] - offset_batch.unsqueeze(-2)
return (
traj_past,
mask_past,
traj_fut,
mask_fut,
mask_loss_batch,
lane_batch,
mask_lane_batch,
offset_batch,
ego_past,
ego_fut,
)
class WaymoDataloaders:
def __init__(self, cfg: Config) -> None:
self.cfg = cfg
def sample_dataloader(self) -> DataLoader:
"""Setup and return sample DataLoader
Returns:
DataLoader: sample DataLoader
"""
dataset = WaymoDataset(self.cfg, "sample")
sample_loader = DataLoader(
dataset=dataset,
batch_size=self.cfg.batch_size,
shuffle=False,
num_workers=self.cfg.num_workers,
collate_fn=dataset.collate_fn,
drop_last=True,
)
return sample_loader
def val_dataloader(
self, drop_last=True, shuffle=False, input_angle=True
) -> DataLoader:
"""Setup and return validation DataLoader
Returns:
DataLoader: validation DataLoader
"""
dataset = WaymoDataset(self.cfg, "val", input_angle)
val_loader = DataLoader(
dataset=dataset,
batch_size=self.cfg.batch_size,
shuffle=shuffle,
num_workers=self.cfg.num_workers,
collate_fn=dataset.collate_fn,
drop_last=drop_last,
)
torch.cuda.empty_cache()
return val_loader
def train_dataloader(
self, drop_last=True, shuffle=True, input_angle=True
) -> DataLoader:
"""Setup and return training DataLoader
Returns:
DataLoader: training DataLoader
"""
dataset = WaymoDataset(self.cfg, "train", input_angle)
train_loader = DataLoader(
dataset=dataset,
batch_size=self.cfg.batch_size,
shuffle=shuffle,
num_workers=self.cfg.num_workers,
collate_fn=dataset.collate_fn,
drop_last=drop_last,
)
torch.cuda.empty_cache()
return train_loader
def test_dataloader(self) -> DataLoader:
"""Setup and return test DataLoader
Returns:
DataLoader: test DataLoader
"""
raise NotImplementedError("The waymo dataloader cannot load test samples yet.")
@staticmethod
def unnormalize_trajectory(
input: torch.Tensor, offset: torch.Tensor
) -> torch.Tensor:
"""Unnormalize trajectory by adding offset to input
Args:
input : (..., (n_sample), num_steps_future, state_dim) tensor of future
trajectory y
offset : (..., state_dim) tensor of offset to add to y
Returns:
Unnormalized trajectory that has the same size as input
"""
assert offset.ndim == 3
batch_size, num_agents = offset.shape[:2]
offset_state_dim = offset.shape[-1]
assert offset_state_dim <= input.shape[-1]
assert input.shape[0] == batch_size
assert input.shape[1] == num_agents
input_copy = input.clone()
input_copy[..., :offset_state_dim] = input_copy[
..., :offset_state_dim
] + offset[..., : input.shape[-1]].reshape(
[batch_size, num_agents, *([1] * (input.ndim - 3)), offset_state_dim]
)
return input_copy