def test_zarr_split(dmg: LocalDataManager, tmp_path: Path, zarr_dataset: ChunkedDataset) -> None: concat_count = 10 zarr_input_path = dmg.require("single_scene.zarr") zarr_concatenated_path = str(tmp_path / f"{uuid4()}.zarr") zarr_concat([zarr_input_path] * concat_count, zarr_concatenated_path) split_infos = [ { "name": f"{uuid4()}.zarr", "split_size_GB": 0.002 }, # cut around 2MB { "name": f"{uuid4()}.zarr", "split_size_GB": 0.001 }, # cut around 0.5MB { "name": f"{uuid4()}.zarr", "split_size_GB": -1 }, ] # everything else scene_splits = zarr_split(zarr_concatenated_path, str(tmp_path), split_infos) # load the zarrs and check elements zarr_concatenated = ChunkedDataset(zarr_concatenated_path) zarr_concatenated.open() for scene_split, split_info in zip(scene_splits, split_infos): zarr_out = ChunkedDataset(str(tmp_path / str(split_info["name"]))) zarr_out.open() # compare elements at the start and end of each scene in both zarrs for idx_scene in range(len(zarr_out.scenes)): # compare elements in the scene input_scene = zarr_concatenated.scenes[scene_split[0] + idx_scene] input_frames = zarr_concatenated.frames[ get_frames_slice_from_scenes(input_scene)] input_agents = zarr_concatenated.agents[ get_agents_slice_from_frames(*input_frames[[0, -1]])] input_tl_faces = zarr_concatenated.tl_faces[ get_tl_faces_slice_from_frames(*input_frames[[0, -1]])] output_scene = zarr_out.scenes[idx_scene] output_frames = zarr_out.frames[get_frames_slice_from_scenes( output_scene)] output_agents = zarr_out.agents[get_agents_slice_from_frames( *output_frames[[0, -1]])] output_tl_faces = zarr_out.tl_faces[get_tl_faces_slice_from_frames( *output_frames[[0, -1]])] assert np.all(input_frames["ego_translation"] == output_frames["ego_translation"]) assert np.all( input_frames["ego_rotation"] == output_frames["ego_rotation"]) assert np.all(input_agents == output_agents) assert np.all(input_tl_faces == output_tl_faces)
def test_get_frames_slice_from_scenes(zarr_dataset: ChunkedDataset) -> None: scene_a = zarr_dataset.scenes[0] frame_slice = get_frames_slice_from_scenes(scene_a) assert len(zarr_dataset.frames) == len(zarr_dataset.frames[frame_slice]) # test e2e starting from scene frame_range = get_frames_slice_from_scenes(zarr_dataset.scenes[0]) agents_range = get_agents_slice_from_frames( *zarr_dataset.frames[frame_range][[0, -1]]) tl_faces_range = get_tl_faces_slice_from_frames( *zarr_dataset.frames[frame_range][[0, -1]]) agents = zarr_dataset.agents[agents_range] tl_faces = zarr_dataset.tl_faces[tl_faces_range] assert len(agents) == len(zarr_dataset.agents) assert len(tl_faces) == len(zarr_dataset.tl_faces)
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 test_zarr_scenes_chunk(dmg: LocalDataManager, tmp_path: Path, zarr_dataset: ChunkedDataset, num_frames_to_copy: int) -> None: # first let's concat so we have multiple scenes concat_count = 10 zarr_input_path = dmg.require("single_scene.zarr") zarr_concatenated_path = str(tmp_path / f"{uuid4()}.zarr") zarr_concat([zarr_input_path] * concat_count, zarr_concatenated_path) # now let's chunk it zarr_chopped_path = str(tmp_path / f"{uuid4()}.zarr") zarr_scenes_chop(zarr_concatenated_path, zarr_chopped_path, num_frames_to_copy=num_frames_to_copy) # open both and compare zarr_concatenated = ChunkedDataset(zarr_concatenated_path) zarr_concatenated.open() zarr_chopped = ChunkedDataset(zarr_chopped_path) zarr_chopped.open() assert len(zarr_concatenated.scenes) == len(zarr_chopped.scenes) assert len( zarr_chopped.frames) == num_frames_to_copy * len(zarr_chopped.scenes) for idx in range(len(zarr_concatenated.scenes)): scene_cat = zarr_concatenated.scenes[idx] scene_chopped = zarr_chopped.scenes[idx] frames_cat = zarr_concatenated.frames[ scene_cat["frame_index_interval"][0]: scene_cat["frame_index_interval"][0] + num_frames_to_copy] frames_chopped = zarr_chopped.frames[get_frames_slice_from_scenes( scene_chopped)] agents_cat = zarr_concatenated.agents[get_agents_slice_from_frames( *frames_cat[[0, -1]])] tl_faces_cat = zarr_concatenated.tl_faces[ get_tl_faces_slice_from_frames(*frames_cat[[0, -1]])] agents_chopped = zarr_chopped.agents[get_agents_slice_from_frames( *frames_chopped[[0, -1]])] tl_faces_chopped = zarr_chopped.tl_faces[ get_tl_faces_slice_from_frames(*frames_chopped[[0, -1]])] assert scene_chopped["host"] == scene_cat["host"] assert scene_chopped["start_time"] == scene_cat["start_time"] assert scene_chopped["end_time"] == scene_cat["end_time"] assert len(frames_chopped) == num_frames_to_copy assert np.all(frames_chopped["ego_translation"] == frames_cat["ego_translation"][:num_frames_to_copy]) assert np.all(frames_chopped["ego_rotation"] == frames_cat["ego_rotation"][:num_frames_to_copy]) assert np.all(agents_chopped == agents_cat) assert np.all(tl_faces_chopped == tl_faces_cat)
def get_frame(self, 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 Args: 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: dict: 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 """ frames = self.dataset.frames[get_frames_slice_from_scenes( self.dataset.scenes[scene_index])] data = self.sample_function(state_index, frames, self.dataset.agents, self.dataset.tl_faces, track_id) # 0,1,C -> C,0,1 image = data["image"].transpose(2, 0, 1) target_positions = np.array(data["target_positions"], dtype=np.float32) # target_velocities = np.array(data["target_velocities"], dtype=np.float32) # target_accelerations = np.array(data["target_accelerations"], dtype=np.float32) target_yaws = np.array(data["target_yaws"], dtype=np.float32) history_positions = np.array(data["history_positions"], dtype=np.float32) # history_velocities = np.array(data["history_velocities"], dtype=np.float32) # history_accelerations = np.array(data["history_accelerations"], dtype=np.float32) # history_yaws = np.array(data["history_yaws"], dtype=np.float32) estimated_positions = np.array(data["estimated_future_positions"], dtype=np.float32) timestamp = frames[state_index]["timestamp"] track_id = np.int64( -1 if track_id is None else track_id) # always a number to avoid crashing torch return { "image": image, "target_positions": target_positions, "target_yaws": target_yaws, "target_availabilities": data["target_availabilities"], "history_positions": history_positions, "history_availabilities": data["history_availabilities"], "estimated_future_positions": estimated_positions, "world_to_image": data["world_to_image"], "track_id": track_id, "timestamp": timestamp, "centroid": data["centroid"], "yaw": data["yaw"], "extent": data["extent"], }
def get_frame(self, 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 Args: 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: dict: 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 """ frames = self.dataset.frames[get_frames_slice_from_scenes(self.dataset.scenes[scene_index])] tl_faces = self.dataset.tl_faces try: if self.cfg["raster_params"]["disable_traffic_light_faces"]: tl_faces = np.empty(0, dtype=self.dataset.tl_faces.dtype) # completely disable traffic light faces except KeyError: warnings.warn( "disable_traffic_light_faces not found in config, this will raise an error in the future", RuntimeWarning, stacklevel=2, ) data = self.sample_function(state_index, frames, self.dataset.agents, tl_faces, track_id) graph = data["image"] target_positions = np.array(data["target_positions"], dtype=np.float32) target_yaws = np.array(data["target_yaws"], dtype=np.float32) history_positions = np.array(data["history_positions"], dtype=np.float32) history_yaws = np.array(data["history_yaws"], dtype=np.float32) timestamp = frames[state_index]["timestamp"] track_id = np.int64(-1 if track_id is None else track_id) # always a number to avoid crashing torch return { "graph": graph, "target_positions": target_positions, "target_yaws": target_yaws, "target_availabilities": data["target_availabilities"], "history_positions": history_positions, "history_yaws": history_yaws, "history_availabilities": data["history_availabilities"], "raster_from_world": data["raster_from_world"], "raster_from_agent": data["raster_from_agent"], "agent_from_world": data["agent_from_world"], "world_from_agent": data["world_from_agent"], "track_id": track_id, "timestamp": timestamp, "centroid": data["centroid"], "yaw": data["yaw"], "extent": data["extent"], "scene_index": scene_index }
def test_get_scene_indices_ego(scene_idx: int, zarr_dataset: ChunkedDataset, dmg: LocalDataManager, cfg: dict) -> None: cfg["raster_params"]["map_type"] = "box_debug" rasterizer = build_rasterizer(cfg, dmg) dataset = EgoDataset(cfg, zarr_dataset, rasterizer) scene_indices = dataset.get_scene_indices(scene_idx) frame_slice = get_frames_slice_from_scenes(zarr_dataset.scenes[scene_idx]) assert scene_indices[0] == frame_slice.start assert scene_indices[-1] == frame_slice.stop - 1
def get_frame_custom(self, scene_index: int, state_index: int, track_id: Optional[int] = None) -> dict: """Customized `get_frame` function, which returns all `data` entries of `sample_function`. A utility function to get the rasterisation and trajectory target for a given agent in a given frame Args: self: Ego Dataset 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: dict: 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 """ frames = self.dataset.frames[get_frames_slice_from_scenes( self.dataset.scenes[scene_index])] tl_faces = self.dataset.tl_faces try: if self.cfg["raster_params"]["disable_traffic_light_faces"]: tl_faces = np.empty(0, dtype=self.dataset.tl_faces.dtype ) # completely disable traffic light faces except KeyError: warnings.warn( "disable_traffic_light_faces not found in config, this will raise an error in the future", RuntimeWarning, stacklevel=2, ) data = self.sample_function(state_index, frames, self.dataset.agents, tl_faces, track_id) # 0,1,C -> C,0,1 image = data["image"].transpose(2, 0, 1) target_positions = np.array(data["target_positions"], dtype=np.float32) target_yaws = np.array(data["target_yaws"], dtype=np.float32) history_positions = np.array(data["history_positions"], dtype=np.float32) history_yaws = np.array(data["history_yaws"], dtype=np.float32) timestamp = frames[state_index]["timestamp"] track_id = np.int64(-1 if track_id is None else track_id) # always a number to avoid crashing torch data["image"] = image data["target_positions"] = target_positions data["target_yaws"] = target_yaws data["history_positions"] = history_positions data["history_yaws"] = history_yaws data["track_id"] = track_id data["timestamp"] = timestamp return data
def test_get_agent_context(zarr_dataset: ChunkedDataset, state_index: int, history_steps: int, future_steps: int) -> None: scene = zarr_dataset.scenes[0] frames = zarr_dataset.frames[get_frames_slice_from_scenes(scene)] agents = zarr_dataset.agents[get_agents_slice_from_frames( *frames[[0, -1]])] tls = zarr_dataset.tl_faces[get_tl_faces_slice_from_frames( *frames[[0, -1]])] frames_his_f, frames_fut_f, agents_his_f, agents_fut_f, tls_his_f, tls_fut_f = get_agent_context( state_index, frames, agents, tls, history_steps, future_steps) # test future using timestamp first_idx = state_index + 1 last_idx = state_index + 1 + future_steps frames_fut = frames[first_idx:last_idx] agents_fut = filter_agents_by_frames(frames_fut, zarr_dataset.agents) tls_fut = filter_tl_faces_by_frames(frames_fut, zarr_dataset.tl_faces) assert np.all(frames_fut_f["timestamp"] == frames_fut["timestamp"]) assert len(agents_fut) == len(agents_fut_f) for idx in range(len(agents_fut)): assert np.all(agents_fut_f[idx] == agents_fut[idx]) assert len(tls_fut) == len(tls_fut_f) for idx in range(len(tls_fut)): assert np.all(tls_fut_f[idx] == tls_fut[idx]) # test past (which is reversed and include present) first_idx = max(state_index - history_steps, 0) last_idx = state_index + 1 frames_his = frames[first_idx:last_idx] agents_his = filter_agents_by_frames(frames_his, zarr_dataset.agents) tls_his = filter_tl_faces_by_frames(frames_his, zarr_dataset.tl_faces) assert np.all(frames_his_f["timestamp"] == frames_his["timestamp"][::-1]) assert len(agents_his) == len(agents_his_f) for idx in range(len(agents_his)): assert np.all(agents_his_f[idx] == agents_his[len(agents_his) - idx - 1]) assert len(tls_his) == len(tls_his_f) for idx in range(len(tls_his)): assert np.all(tls_his_f[idx] == tls_his[len(tls_his) - idx - 1])
def get_frame(self, scene_index: int, state_index: int) -> dict: """ A utility function to get the rasterisation and trajectory target for a given agent in a given frame Args: 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: dict: 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 """ frames_slice = get_frames_slice_from_scenes( self.zarr_root[SCENE_ARRAY_KEY][scene_index]) frames = self.zarr_root[FRAME_ARRAY_KEY][frames_slice] timestamp = frames[state_index]["timestamp"] data = self.sample_function(state_index, frames) results = { "timestamp": timestamp, } if self.with_history: results.update({ "ego_centroid": data["ego_centroid"], "ego_speed": data["ego_speed"], "history_tl_faces": data["history_tl_faces"], "history_agents": data["history_agents"], }) else: results.update(({ "ego_centroid": data["ego_centroid"], "ego_speed": data["ego_speed"], "ego_yaw": data["ego_yaw"], "tl_faces": data["tl_faces"], "agents": data["agents"], })) if self.return_indices: results.update({ "scene_index": scene_index, "state_index": state_index }) return results
def test_unroll(zarr_cat_dataset: ChunkedDataset, dmg: LocalDataManager, cfg: dict) -> None: rasterizer = build_rasterizer(cfg, dmg) # change the first yaw of scene 1 # this will break if some broadcasting happens zarr_cat_dataset.frames = np.asarray(zarr_cat_dataset.frames) slice_frames = get_frames_slice_from_scenes(zarr_cat_dataset.scenes[1]) rot = zarr_cat_dataset.frames[slice_frames.start]["ego_rotation"].copy() zarr_cat_dataset.frames[ slice_frames.start]["ego_rotation"] = yaw_as_rotation33( rotation33_as_yaw(rot + 0.75)) scene_indices = list(range(len(zarr_cat_dataset.scenes))) ego_dataset = EgoDataset(cfg, zarr_cat_dataset, rasterizer) # control only agents at T0, control them forever sim_cfg = SimulationConfig(use_ego_gt=False, use_agents_gt=False, disable_new_agents=True, distance_th_close=1000, distance_th_far=1000, num_simulation_steps=10) # ego will move by 1 each time ego_model = MockModel(advance_x=1.0) # agents will move by 0.5 each time agents_model = MockModel(advance_x=0.5) sim = ClosedLoopSimulator(sim_cfg, ego_dataset, torch.device("cpu"), ego_model, agents_model) sim_outputs = sim.unroll(scene_indices) # check ego movement for sim_output in sim_outputs: ego_tr = sim_output.simulated_ego[ "ego_translation"][:sim_cfg.num_simulation_steps, :2] ego_dist = np.linalg.norm(np.diff(ego_tr, axis=0), axis=-1) assert np.allclose(ego_dist, 1.0) ego_tr = sim_output.simulated_ego_states[:sim_cfg.num_simulation_steps, TrajectoryStateIndices. X:TrajectoryStateIndices.Y + 1] ego_dist = np.linalg.norm(np.diff(ego_tr.numpy(), axis=0), axis=-1) assert np.allclose(ego_dist, 1.0, atol=1e-3) # all rotations should be the same as the first one as the MockModel outputs 0 for that rots_sim = sim_output.simulated_ego[ "ego_rotation"][:sim_cfg.num_simulation_steps] r_rep = sim_output.recorded_ego["ego_rotation"][0] for r_sim in rots_sim: assert np.allclose(rotation33_as_yaw(r_sim), rotation33_as_yaw(r_rep), atol=1e-2) # all rotations should be the same as the first one as the MockModel outputs 0 for that rots_sim = sim_output.simulated_ego_states[:sim_cfg. num_simulation_steps, TrajectoryStateIndices. THETA] r_rep = sim_output.recorded_ego_states[0, TrajectoryStateIndices.THETA] for r_sim in rots_sim: assert np.allclose(r_sim, r_rep, atol=1e-2) # check agents movements for sim_output in sim_outputs: # we need to know which agents were controlled during simulation # TODO: this is not ideal, we should keep track of them through the simulation sim_dataset = SimulationDataset.from_dataset_indices( ego_dataset, [sim_output.scene_id], sim_cfg) sim_dataset.rasterise_agents_frame_batch( 0) # this will fill agents_tracked agents_tracks = [el[1] for el in sim_dataset._agents_tracked] for track_id in agents_tracks: states = sim_output.simulated_agents agents = filter_agents_by_track_id( states, track_id)[:sim_cfg.num_simulation_steps] agent_dist = np.linalg.norm(np.diff(agents["centroid"], axis=0), axis=-1) assert np.allclose(agent_dist, 0.5)