l5kit.dataset package

class l5kit.dataset.AgentDataset(cfg: dict, zarr_dataset: l5kit.data.zarr_dataset.ChunkedStateDataset, rasterizer: l5kit.rasterization.rasterizer.Rasterizer, perturbation: Optional[l5kit.kinematic.perturbation.Perturbation] = None, agents_mask: Optional[numpy.ndarray] = None)

Bases: l5kit.dataset.ego.EgoDataset

get_frame_indices(frame_idx: int) → numpy.ndarray

Get indices for the given frame. Here __getitem__ iterate over valid agents indices. This means __getitem__(0) matches the first valid agent in the dataset. :param frame_idx: index of the scene :type frame_idx: int

Returns

indices that can be used for indexing with __getitem__

Return type

np.ndarray

get_scene_dataset(scene_index: int) → l5kit.dataset.agent.AgentDataset

Differs from parent only in the return type. Instead of doing everything from scratch, we rely on super call and fix the agents_mask

get_scene_indices(scene_idx: int) → numpy.ndarray

Get indices for the given scene. Here __getitem__ iterate over valid agents indices. This means __getitem__(0) matches the first valid agent in the dataset. :param scene_idx: index of the scene :type scene_idx: int

Returns

indices that can be used for indexing with __getitem__

Return type

np.ndarray

load_agents_mask() → numpy.ndarray

Loads a boolean mask of the agent availability stored into the zarr. Performs some sanity check against cfg. Returns: a boolean mask of the same length of the dataset agents

class l5kit.dataset.EgoDataset(cfg: dict, zarr_dataset: l5kit.data.zarr_dataset.ChunkedStateDataset, rasterizer: l5kit.rasterization.rasterizer.Rasterizer, perturbation: Optional[l5kit.kinematic.perturbation.Perturbation] = None)

Bases: torch.utils.data.dataset.Dataset

get_frame(scene_index: int, state_index: int, track_id: Optional[int] = None) → dict

A utility function to get the rasterisation and trajectory target for a given agent in a given frame

Parameters

  • scene_index (int) – the index of the scene in the zarr

  • state_index (int) – a relative frame index in the scene

  • track_id (Optional[int]) – the agent to rasterize or None for the AV

Returns

the rasterised image, the target trajectory (position and yaw) along with their availability, the 2D matrix to center that agent, the agent track (-1 if ego) and the timestamp

Return type

dict

get_frame_indices(frame_idx: int) → numpy.ndarray

Get indices for the given frame. EgoDataset iterates over frames, so this will be a single element :param frame_idx: index of the scene :type frame_idx: int

Returns

indices that can be used for indexing with __getitem__

Return type

np.ndarray

get_scene_dataset(scene_index: int) → l5kit.dataset.ego.EgoDataset

Returns another EgoDataset dataset where the underlying data can be modified. This is possible because, even if it supports the same interface, this dataset is np.ndarray based.

Parameters

scene_index (int) – the scene index of the new dataset

Returns

A valid EgoDataset dataset with a copy of the data

Return type

EgoDataset

get_scene_indices(scene_idx: int) → numpy.ndarray

Get indices for the given scene. EgoDataset iterates over frames, so this is just a matter of finding the scene boundaries. :param scene_idx: index of the scene :type scene_idx: int

Returns

indices that can be used for indexing with __getitem__

Return type

np.ndarray

l5kit.dataset.build_dataloader(cfg: Dict, split: str, data_manager: l5kit.data.local_data_manager.DataManager, dataset_class: Callable, rasterizer: l5kit.rasterization.rasterizer.Rasterizer, perturbation: Optional[l5kit.kinematic.perturbation.Perturbation] = None) → torch.utils.data.dataloader.DataLoader

Function to build a dataloader from a dataset of dataset_class. Note we have to pass rasterizer and perturbation as the factory functions for those are likely to change between repos.

Parameters

  • cfg (dict) – configuration dict

  • split (str) – this will be used to index the cfg to get the correct datasets (train or val currently)

  • data_manager (DataManager) – manager for resolving paths

  • dataset_class (Callable) – a class object (EgoDataset or AgentDataset currently) to build the dataset

  • rasterizer (Rasterizer) – the rasterizer for the dataset

  • perturbation (Optional[Perturbation]) – an optional perturbation object

Returns

pytorch Dataloader object built with Concat and Sub datasets

Return type

DataLoader