def test_simulation_dataset_build(zarr_cat_dataset: ChunkedDataset,
dmg: LocalDataManager, cfg: dict,
tmp_path: Path) -> None:
# modify one frame to ensure everything works also when scenes are different
zarr_cat_dataset.frames = np.asarray(zarr_cat_dataset.frames)
for scene_idx in range(len(zarr_cat_dataset.scenes)):
frame_slice = get_frames_slice_from_scenes(zarr_cat_dataset.scenes)
zarr_cat_dataset.frames[
frame_slice.start]["ego_translation"] += np.random.randn(3)
rasterizer = build_rasterizer(cfg, dmg)
ego_dataset = EgoDataset(cfg, zarr_cat_dataset, rasterizer)
sim_cfg = SimulationConfig(use_ego_gt=True,
use_agents_gt=True,
disable_new_agents=False,
distance_th_far=30,
distance_th_close=10)
# we should be able to create the same object by using both constructor and factory
scene_indices = list(range(len(zarr_cat_dataset.scenes)))
scene_dataset_batch: Dict[int, EgoDataset] = {}
for scene_idx in scene_indices:
scene_dataset = ego_dataset.get_scene_dataset(scene_idx)
scene_dataset_batch[scene_idx] = scene_dataset
sim_1 = SimulationDataset(scene_dataset_batch, sim_cfg)
sim_2 = SimulationDataset.from_dataset_indices(ego_dataset, scene_indices,
sim_cfg)
for (k_1, v_1), (k_2, v_2) in zip(sim_1.scene_dataset_batch.items(),
sim_2.scene_dataset_batch.items()):
assert k_1 == k_2
assert np.allclose(v_1.dataset.frames["ego_translation"],
v_2.dataset.frames["ego_translation"])
def from_dataset_indices(dataset: EgoDataset, scene_indices: List[int],
sim_cfg: SimulationConfig) -> "SimulationDataset":
"""Create a SimulationDataset by picking indices from the provided dataset
:param dataset: the EgoDataset
:param scene_indices: scenes from the EgoDataset to pick
:param sim_cfg: a simulation config
:return: the new SimulationDataset
"""
if len(np.unique(scene_indices)) != len(scene_indices):
raise ValueError(f"can't simulate repeated scenes: {scene_indices}")
if np.any(np.asarray(scene_indices) >= len(dataset.dataset.scenes)):
raise ValueError(
f"can't pick indices {scene_indices} from dataset with length: {len(dataset.dataset.scenes)}")
scene_dataset_batch: Dict[int, EgoDataset] = {} # dicts preserve insertion order
for scene_idx in scene_indices:
scene_dataset = dataset.get_scene_dataset(scene_idx)
scene_dataset_batch[scene_idx] = scene_dataset
return SimulationDataset(scene_dataset_batch, sim_cfg)
