def test_dataset_frames_subset_invalid(zarr_dataset: ChunkedDataset) -> None:
with pytest.raises(ValueError):
get_frames_subset(zarr_dataset, 0, 10) # not in numpy
with pytest.raises(ValueError):
get_frames_subset(zarr_dataset, -1, 10) # invalid start
with pytest.raises(ValueError):
get_frames_subset(
zarr_dataset, len(zarr_dataset.frames), len(zarr_dataset.frames) + 1
) # invalid start
with pytest.raises(ValueError):
get_frames_subset(zarr_dataset, 10, 9) # invalid end
with pytest.raises(ValueError):
get_frames_subset(zarr_dataset, 10, len(zarr_dataset.frames) + 1) # invalid end
def __init__(self, scene_dataset_batch: Dict[int, EgoDataset], sim_cfg: SimulationConfig) -> None:
"""This class allows to:
- rasterise the same frame across multiple scenes for ego;
- rasterise the same frame across multiple scenes for multiple agents;
- filter agents based on distance to ego;
- set ego in future frames;
- set agents in future frames;
.. note:: only vehicles (car label) are picked as agents
:param scene_dataset_batch: a mapping from scene index to EgoDataset
:param sim_cfg: the simulation config
"""
if not len(scene_dataset_batch):
raise ValueError("can't build a simulation dataset with an empty batch")
self.scene_dataset_batch: Dict[int, EgoDataset] = scene_dataset_batch
self.sim_cfg = sim_cfg
# we must limit the scenes to the part which will be simulated
# we cut each scene so that it starts from there and ends after `num_simulation_steps`
start_frame_idx = self.sim_cfg.start_frame_index
if self.sim_cfg.num_simulation_steps is None:
end_frame_idx = self.get_min_len()
else:
end_frame_idx = start_frame_idx + self.sim_cfg.num_simulation_steps
if end_frame_idx > self.get_min_len():
raise ValueError(f"can't unroll until frame {end_frame_idx}, length is {self.get_min_len()}")
for scene_idx in scene_dataset_batch:
zarr_dt = self.scene_dataset_batch[scene_idx].dataset
self.scene_dataset_batch[scene_idx].dataset = get_frames_subset(zarr_dt, start_frame_idx, end_frame_idx)
# this is the only stateful field we need to change for EgoDataset, it's used in bisect
frame_index_ends = self.scene_dataset_batch[scene_idx].dataset.scenes["frame_index_interval"][:, 1]
self.scene_dataset_batch[scene_idx].cumulative_sizes = frame_index_ends
# buffer used to keep track of tracked agents during unroll as tuples of scene_idx, agent_idx
self._agents_tracked: Set[Tuple[int, int]] = set()
if self.sim_cfg.disable_new_agents:
# we disable all agents that wouldn't be picked at frame 0
for scene_idx, dt_ego in self.scene_dataset_batch.items():
dataset_zarr = dt_ego.dataset
frame = dataset_zarr.frames[0]
ego_pos = frame["ego_translation"][:2]
agents = dataset_zarr.agents
frame_agents = filter_agents_by_frames(frame, agents)[0]
frame_agents = self._filter_agents(scene_idx, frame_agents, ego_pos)
disable_agents(dataset_zarr, allowlist=frame_agents["track_id"])
# keep track of original dataset
self.recorded_scene_dataset_batch = deepcopy(self.scene_dataset_batch)
def test_mock_dataset_frames_subset() -> None:
zarr_dataset = ChunkedDataset("")
zarr_dataset.scenes = np.zeros(1, dtype=SCENE_DTYPE)
zarr_dataset.scenes[0]["frame_index_interval"] = (0, 4)
zarr_dataset.frames = np.zeros(4, dtype=FRAME_DTYPE)
zarr_dataset.frames["agent_index_interval"] = [(0, 1), (1, 2), (2, 3), (3, 4)]
zarr_dataset.agents = np.zeros(4, dtype=AGENT_DTYPE)
zarr_dataset.agents["track_id"] = np.arange(4)
zarr_dataset.tl_faces = np.zeros(0, dtype=TL_FACE_DTYPE)
frame_start = 1
frame_end = 3
zarr_cut = get_frames_subset(zarr_dataset, frame_start, frame_end)
assert np.all(zarr_cut.agents["track_id"] == [1, 2])
frame_start = 0
frame_end = 3
zarr_cut = get_frames_subset(zarr_dataset, frame_start, frame_end)
assert np.all(zarr_cut.agents["track_id"] == [0, 1, 2])
frame_start = 2
frame_end = 4
zarr_cut = get_frames_subset(zarr_dataset, frame_start, frame_end)
assert np.all(zarr_cut.agents["track_id"] == [2, 3])
def test_dataset_frames_subset(zarr_dataset: ChunkedDataset) -> None:
zarr_dataset = zarr_dataset.get_scene_dataset(0)
frame_start = 10
frame_end = 25
zarr_cut = get_frames_subset(zarr_dataset, frame_start, frame_end)
assert len(zarr_cut.scenes) == 1
assert len(zarr_cut.frames) == frame_end - frame_start
assert np.all(
zarr_cut.frames["ego_translation"]
== zarr_dataset.frames["ego_translation"][frame_start:frame_end]
)
agents_slice = get_agents_slice_from_frames(
*zarr_dataset.frames[[frame_start, frame_end - 1]]
)
tls_slice = get_tl_faces_slice_from_frames(
*zarr_dataset.frames[[frame_start, frame_end - 1]]
)
assert np.all(zarr_cut.agents == zarr_dataset.agents[agents_slice])
assert np.all(zarr_cut.tl_faces == zarr_dataset.tl_faces[tls_slice])
assert np.all(zarr_cut.scenes["frame_index_interval"] == (0, len(zarr_cut.frames)))