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 _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 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])
