def test_compute_mse_error(tmp_path: Path) -> None:
data = ChunkedDataset(path="./l5kit/tests/artefacts/single_scene.zarr")
data.open()
export_zarr_to_ground_truth_csv(data, str(tmp_path / "gt1.csv"), 0, 50, 0.5)
data.open() # avoid double select_agents
export_zarr_to_ground_truth_csv(data, str(tmp_path / "gt2.csv"), 0, 50, 0.5)
err = compute_mse_error_csv(str(tmp_path / "gt1.csv"), str(tmp_path / "gt2.csv"))
assert np.all(err == 0.0)
data_fake = ChunkedDataset(str(tmp_path))
data_fake.scenes = np.asarray(data.scenes).copy()
data_fake.frames = np.asarray(data.frames).copy()
data_fake.agents = np.asarray(data.agents).copy()
data_fake.root = data.root
data_fake.agents["centroid"] += np.random.rand(*data_fake.agents["centroid"].shape)
export_zarr_to_ground_truth_csv(data_fake, str(tmp_path / "gt3.csv"), 0, 50, 0.5)
err = compute_mse_error_csv(str(tmp_path / "gt1.csv"), str(tmp_path / "gt3.csv"))
assert np.any(err > 0.0)
# test invalid conf by removing lines in gt1
with open(str(tmp_path / "gt4.csv"), "w") as fp:
lines = open(str(tmp_path / "gt1.csv")).readlines()
fp.writelines(lines[:-10])
with pytest.raises(ValueError):
compute_mse_error_csv(str(tmp_path / "gt1.csv"), str(tmp_path / "gt4.csv"))
def test_compute_mse_error(tmp_path: Path,
zarr_dataset: ChunkedDataset) -> None:
export_zarr_to_ground_truth_csv(zarr_dataset, str(tmp_path / "gt1.csv"),
10, 50, 0.5)
export_zarr_to_ground_truth_csv(zarr_dataset, str(tmp_path / "gt2.csv"),
10, 50, 0.5)
err = compute_mse_error_csv(str(tmp_path / "gt1.csv"),
str(tmp_path / "gt2.csv"))
assert np.all(err == 0.0)
data_fake = ChunkedDataset(str(tmp_path))
data_fake.scenes = np.asarray(zarr_dataset.scenes).copy()
data_fake.frames = np.asarray(zarr_dataset.frames).copy()
data_fake.agents = np.asarray(zarr_dataset.agents).copy()
data_fake.agents["centroid"] += np.random.rand(
*data_fake.agents["centroid"].shape) * 1e-2
export_zarr_to_ground_truth_csv(data_fake, str(tmp_path / "gt3.csv"), 10,
50, 0.5)
err = compute_mse_error_csv(str(tmp_path / "gt1.csv"),
str(tmp_path / "gt3.csv"))
assert np.any(err > 0.0)
# test invalid conf by removing lines in gt1
with open(str(tmp_path / "gt4.csv"), "w") as fp:
lines = open(str(tmp_path / "gt1.csv")).readlines()
fp.writelines(lines[:-10])
with pytest.raises(ValueError):
compute_mse_error_csv(str(tmp_path / "gt1.csv"),
str(tmp_path / "gt4.csv"))
def _get_simple_dataset(self) -> ChunkedDataset:
# build a simple dataset with 3 frames
# frame 0:
# agent 0
# agent 1
# agent 2
# frame 1:
# agent 0
# agent 1
# frame 2:
# agent 0
dataset = ChunkedDataset("")
dataset.scenes = np.zeros(1, dtype=SCENE_DTYPE)
dataset.frames = np.zeros(3, dtype=FRAME_DTYPE)
dataset.agents = np.zeros(6, dtype=AGENT_DTYPE)
dataset.scenes[0]["frame_index_interval"] = (0, 3)
dataset.frames["agent_index_interval"] = [(0, 3), (3, 5), (5, 6)]
dataset.agents["track_id"] = [0, 1, 2, 0, 1, 0]
# set properties to something different than 0
dataset.agents["centroid"] = np.random.rand(*dataset.agents["centroid"].shape)
dataset.agents["yaw"] = np.random.rand(*dataset.agents["yaw"].shape)
dataset.agents["extent"] = np.random.rand(*dataset.agents["extent"].shape)
return dataset
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 get_frames_subset(dataset: ChunkedDataset, frame_start_idx: int,
frame_end_idx: int) -> ChunkedDataset:
"""Get a new dataset with frames between start (included) and end (excluded).
Assumptions:
- the dataset has only 1 scene
- the dataset is in numpy format and not zarr anymore
:param dataset: the single-scene dataset.
:param frame_start_idx: first frame to keep.
:param frame_end_idx: where to stop taking frames (excluded).
"""
if not len(dataset.scenes) == 1:
raise ValueError(
f"dataset should have a single scene, got {len(dataset.scenes)}")
if not isinstance(dataset.agents, np.ndarray):
raise ValueError("dataset agents should be an editable np array")
if not isinstance(dataset.tl_faces, np.ndarray):
raise ValueError("dataset tls should be an editable np array")
if not isinstance(dataset.frames, np.ndarray):
raise ValueError("dataset frames should be an editable np array")
if frame_start_idx >= len(dataset.frames):
raise ValueError(
f"frame start {frame_start_idx} is over the length of the dataset")
if frame_end_idx > len(dataset.frames):
raise ValueError(
f"frame end {frame_end_idx} is over the length of the dataset")
if frame_start_idx >= frame_end_idx:
raise ValueError(
f"end frame {frame_end_idx} should be higher than start {frame_start_idx}"
)
if frame_start_idx < 0:
raise ValueError(f"start frame {frame_start_idx} should be positive")
new_dataset = ChunkedDataset("")
new_dataset.scenes = dataset.scenes.copy()
new_dataset.scenes[0]["start_time"] = dataset.frames[frame_start_idx][
"timestamp"]
new_dataset.scenes[0]["end_time"] = dataset.frames[frame_end_idx -
1]["timestamp"]
new_dataset.frames = dataset.frames[frame_start_idx:frame_end_idx].copy()
new_dataset.scenes[0]["frame_index_interval"] = (0,
len(new_dataset.frames))
agent_slice = get_agents_slice_from_frames(
*dataset.frames[[frame_start_idx, frame_end_idx - 1]])
tls_slice = get_tl_faces_slice_from_frames(
*dataset.frames[[frame_start_idx, frame_end_idx - 1]])
new_dataset.frames["agent_index_interval"] -= new_dataset.frames[
"agent_index_interval"][0, 0]
new_dataset.frames[
"traffic_light_faces_index_interval"] -= new_dataset.frames[
"traffic_light_faces_index_interval"][0, 0]
new_dataset.agents = dataset.agents[agent_slice].copy()
new_dataset.tl_faces = dataset.tl_faces[tls_slice].copy()
return new_dataset
def dataset(tmp_path: Path) -> ChunkedDataset:
dataset = ChunkedDataset(str(tmp_path))
dataset.scenes = np.zeros(1, dtype=dataset.scenes.dtype)
dataset.frames = np.zeros(SCENE_LENGTH, dtype=dataset.frames.dtype)
dataset.agents = np.zeros(SCENE_LENGTH, dtype=dataset.agents.dtype)
dataset.scenes[0]["frame_index_interval"] = (0, SCENE_LENGTH)
for idx in range(len(dataset.frames)):
dataset.frames[idx]["agent_index_interval"] = (idx, idx + 1)
dataset.frames[idx]["timestamp"] = idx
for idx in range(len(dataset.agents)):
# we don't check moving anymore, so the agent can stay still
dataset.agents[idx]["extent"] = (5, 5, 5)
dataset.agents[idx]["yaw"] = 0
dataset.agents[idx]["track_id"] = 1
dataset.agents[idx]["label_probabilities"][3] = 1.0
return dataset
def _mock_dataset() -> ChunkedDataset:
zarr_dt = ChunkedDataset("")
zarr_dt.scenes = np.zeros(1, dtype=SCENE_DTYPE)
zarr_dt.scenes["frame_index_interval"][0] = (0, 4)
zarr_dt.frames = np.zeros(4, dtype=FRAME_DTYPE)
zarr_dt.frames["agent_index_interval"][0] = (0, 3)
zarr_dt.frames["agent_index_interval"][1] = (3, 5)
zarr_dt.frames["agent_index_interval"][2] = (5, 6)
zarr_dt.frames["agent_index_interval"][3] = (6, 6)
zarr_dt.agents = np.zeros(6, dtype=AGENT_DTYPE)
# all agents except the first one are valid
zarr_dt.agents["label_probabilities"][1:, 3] = 1
# FRAME 0
# second agent is close to ego and has id 1
zarr_dt.agents["track_id"][1] = 1
zarr_dt.agents["centroid"][1] = (1, 1)
# third agent is too far and has id 2
zarr_dt.agents["track_id"][2] = 2
zarr_dt.agents["centroid"][2] = (100, 100)
# FRAME 1
# track 1 agent is still close to ego
zarr_dt.agents["track_id"][3] = 1
zarr_dt.agents["centroid"][3] = (1, 2)
# track 2 is now close enough
zarr_dt.agents["track_id"][4] = 2
zarr_dt.agents["centroid"][4] = (1, 1)
# FRAME 2
# track 1 agent is far
zarr_dt.agents["track_id"][5] = 1
zarr_dt.agents["centroid"][5] = (100, 100)
# FRAME 3 is empty
zarr_dt.tl_faces = np.zeros(0, dtype=TL_FACE_DTYPE)
return zarr_dt
def _get_simple_dataset(self) -> ChunkedDataset:
# build a simple dataset with 3 frames
# frame 0:
# agent 0
# agent 1
# agent 2
# frame 1:
# agent 0
# agent 1
# frame 2:
# agent 0
dataset = ChunkedDataset("")
dataset.scenes = np.zeros(1, dtype=SCENE_DTYPE)
dataset.frames = np.zeros(3, dtype=FRAME_DTYPE)
dataset.agents = np.zeros(6, dtype=AGENT_DTYPE)
dataset.scenes[0]["frame_index_interval"] = (0, 3)
dataset.frames["agent_index_interval"] = [(0, 3), (3, 5), (5, 6)]
dataset.agents["track_id"] = [0, 1, 2, 0, 1, 0]
return dataset
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_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)
