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)