def test_get_scene_indices_agent(scene_idx: int, zarr_dataset: ChunkedDataset,
dmg: LocalDataManager, cfg: dict) -> None:
cfg["raster_params"]["map_type"] = "box_debug"
rasterizer = build_rasterizer(cfg, dmg)
dataset = AgentDataset(cfg, zarr_dataset, rasterizer)
# test only first 10 elements
scene_indices = dataset.get_scene_indices(scene_idx)[:10]
agents = np.asarray(dataset.dataset.agents)[dataset.agents_mask][:10]
for agent, idx in zip(agents, scene_indices):
id_agent = dataset[idx]["track_id"]
assert id_agent == agent["track_id"]
dataset = AgentDataset(cfg, zarr_dataset, rast)
data = dataset[80]
im = data["image"].transpose(1, 2, 0)
im = dataset.rasterizer.to_rgb(im)
target_positions_pixels = transform_points(data["target_positions"] + data["centroid"][:2], data["world_to_image"])
draw_trajectory(im, target_positions_pixels, data["target_yaws"], TARGET_POINTS_COLOR)
fig = plt.subplots(figsize=(10,10))
plt.imshow(im[::-1])
plt.show()
dataset = EgoDataset(cfg, zarr_dataset, rast)
scene_idx = 2
indexes = dataset.get_scene_indices(scene_idx)
images = []
for idx in indexes:
data = dataset[idx]
im = data["image"].transpose(1, 2, 0)
im = dataset.rasterizer.to_rgb(im)
target_positions_pixels = transform_points(data["target_positions"] + data["centroid"][:2], data["world_to_image"])
center_in_pixels = np.asarray(cfg["raster_params"]["ego_center"]) * cfg["raster_params"]["raster_size"]
draw_trajectory(im, target_positions_pixels, data["target_yaws"], TARGET_POINTS_COLOR)
clear_output(wait=True)
display(PIL.Image.fromarray(im[::-1]))
#sample submission
sample_submission = pd.read_csv(f'{DIR_PATH}multi_mode_sample_submission.csv')
