l5kit.dataset.ego module

class l5kit.dataset.ego.BaseEgoDataset(cfg: dict, zarr_dataset: l5kit.data.zarr_dataset.ChunkedDataset)

Bases: Generic[torch.utils.data.dataset.T_co]

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 in (Cx0x1) if the rast is not None, 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.BaseEgoDataset

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

class l5kit.dataset.ego.EgoAgentDatasetVectorized(cfg: dict, zarr_dataset: l5kit.data.zarr_dataset.ChunkedDataset, vectorizer: l5kit.vectorization.vectorizer.Vectorizer, perturbation: Optional[l5kit.kinematic.perturbation.Perturbation] = None, agents_mask: Optional[numpy.ndarray] = None, eval_mode: bool = False)

Bases: Generic[torch.utils.data.dataset.T_co]

Get a PyTorch dataset object that can be used to train DNNs with vectorized input. Ego features are added to the agent features to treat ego as an agent.

Parameters
  • cfg (dict) – configuration file

  • zarr_dataset (ChunkedDataset) – the raw zarr dataset

  • vectorizer (Vectorizer) – a object that supports vectorization around an AV

  • perturbation (Optional[Perturbation]) – an object that takes care of applying trajectory perturbations. None if not desired

  • agents_mask (Optional[np.ndarray]) – custom boolean mask of the agent availability.

  • eval_mode (bool) – enable eval mode (iterates over agent, similarly to AgentDataset).

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.EgoAgentDatasetVectorized

Returns another EgoAgentDatasetVectorized 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 EgoAgentDatasetVectorized dataset with a copy of the data

Return type

EgoAgentDatasetVectorized

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

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.ego.EgoDataset(cfg: dict, zarr_dataset: l5kit.data.zarr_dataset.ChunkedDataset, rasterizer: l5kit.rasterization.rasterizer.Rasterizer, perturbation: Optional[l5kit.kinematic.perturbation.Perturbation] = None)

Bases: Generic[torch.utils.data.dataset.T_co]

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 in (Cx0x1) if the rast is not None, 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_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

class l5kit.dataset.ego.EgoDatasetVectorized(cfg: dict, zarr_dataset: l5kit.data.zarr_dataset.ChunkedDataset, vectorizer: l5kit.vectorization.vectorizer.Vectorizer, perturbation: Optional[l5kit.kinematic.perturbation.Perturbation] = None)

Bases: Generic[torch.utils.data.dataset.T_co]

get_scene_dataset(scene_index: int) l5kit.dataset.ego.EgoDatasetVectorized

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