def get_partial(
cfg: dict,
history_num_frames: int,
history_step_size: int,
history_step_time: float,
future_num_frames: int,
future_step_size: int,
future_step_time: float,
) -> Callable:
rast_params = cfg["raster_params"]
render_context = RenderContext(
raster_size_px=np.array(cfg["raster_params"]["raster_size"]),
pixel_size_m=np.array(cfg["raster_params"]["pixel_size"]),
center_in_raster_ratio=np.array(cfg["raster_params"]["ego_center"]),
)
rasterizer = StubRasterizer(render_context, rast_params["filter_agents_threshold"],)
return functools.partial(
generate_agent_sample,
render_context=render_context,
history_num_frames=history_num_frames,
history_step_size=history_step_size,
history_step_time=history_step_time,
future_num_frames=future_num_frames,
future_step_size=future_step_size,
future_step_time=future_step_time,
filter_agents_threshold=rast_params["filter_agents_threshold"],
rasterizer=rasterizer,
)
def build_rasterizer(cfg: dict,
data_manager: DataManager,
debug: bool = False) -> Rasterizer:
raster_cfg = cfg["raster_params"]
map_type = raster_cfg["map_type"]
if map_type == "semantic_graph":
dataset_meta_key = raster_cfg["dataset_meta_key"]
filter_agents_threshold = raster_cfg["filter_agents_threshold"]
history_num_frames = cfg["model_params"]["history_num_frames"]
render_context = RenderContext(
raster_size_px=np.array(raster_cfg["raster_size"]),
pixel_size_m=np.array(raster_cfg["pixel_size"]),
center_in_raster_ratio=np.array(raster_cfg["ego_center"]),
)
semantic_map_filepath = data_manager.require(
raster_cfg["semantic_map_key"])
dataset_meta = _load_metadata(dataset_meta_key, data_manager)
world_to_ecef = np.array(dataset_meta["world_to_ecef"],
dtype=np.float64)
return SemGraphRasterizer(render_context, filter_agents_threshold,
history_num_frames, semantic_map_filepath,
world_to_ecef, debug)
else:
return l5kit_build_rasterizer(cfg, data_manager)
def test_same_displacement(
cfg: dict,
zarr_dataset: ChunkedDataset,
base_displacement: np.ndarray,
raster_size: tuple,
ego_center: tuple,
pixel_size: tuple,
) -> None:
cfg["raster_params"]["raster_size"] = raster_size
cfg["raster_params"]["pixel_size"] = np.asarray(pixel_size)
cfg["raster_params"]["ego_center"] = np.asarray(ego_center)
render_context = RenderContext(
np.asarray(raster_size),
np.asarray(pixel_size),
np.asarray(ego_center),
set_origin_to_bottom=cfg["raster_params"]["set_origin_to_bottom"],
)
dataset = EgoDataset(
cfg,
zarr_dataset,
StubRasterizer(render_context),
)
data = dataset[0]
assert np.allclose(data["target_positions"], base_displacement)
def from_cfg(cfg, data_manager):
raster_cfg = cfg["raster_params"]
# map_type = raster_cfg["map_type"]
dataset_meta_key = raster_cfg["dataset_meta_key"]
render_context = RenderContext(
raster_size_px=np.array(raster_cfg["raster_size"]),
pixel_size_m=np.array(raster_cfg["pixel_size"]),
center_in_raster_ratio=np.array(raster_cfg["ego_center"]),
)
# filter_agents_threshold = raster_cfg["filter_agents_threshold"]
# history_num_frames = cfg["model_params"]["history_num_frames"]
semantic_map_filepath = data_manager.require(
raster_cfg["semantic_map_key"])
try:
dataset_meta = _load_metadata(dataset_meta_key, data_manager)
world_to_ecef = np.array(dataset_meta["world_to_ecef"],
dtype=np.float64)
except (KeyError, FileNotFoundError
): # TODO remove when new dataset version is available
world_to_ecef = get_hardcoded_world_to_ecef()
lane_color_mode = raster_cfg.get("lane_color_mode", "normal")
return ChannelSemanticTLRasterizer(render_context,
semantic_map_filepath,
world_to_ecef,
lane_color_mode=lane_color_mode)
def __init__(
self,
cfg: dict,
zarr_dataset: ChunkedDataset,
rasterizer: Rasterizer,
perturbation: Optional[Perturbation] = None,
min_frame_history: int = 1,
min_frame_future: int = 10,
):
super(MultiAgentDataset, self).__init__(cfg, zarr_dataset, rasterizer, perturbation)
render_context = RenderContext(
raster_size_px=np.array(cfg["raster_params"]["raster_size"]),
pixel_size_m=np.array(cfg["raster_params"]["pixel_size"]),
center_in_raster_ratio=np.array(cfg["raster_params"]["ego_center"]),
)
self.sample_function = partial(
generate_multi_agent_sample,
render_context=render_context,
history_num_frames=cfg["model_params"]["history_num_frames"],
history_step_size=cfg["model_params"]["history_step_size"],
future_num_frames=cfg["model_params"]["future_num_frames"],
future_step_size=cfg["model_params"]["future_step_size"],
filter_agents_threshold=cfg["raster_params"]["filter_agents_threshold"],
rasterizer=rasterizer,
perturbation=perturbation,
min_frame_history=min_frame_history,
min_frame_future=min_frame_future,
)
def test_scene_index_interval(dataset_cls: Callable, scene_idx: int,
zarr_dataset: ChunkedDataset, cfg: dict) -> None:
render_context = RenderContext(np.asarray((100, 100)),
np.asarray((0.25, 0.25)),
np.asarray((0.5, 0.5)))
rasterizer = StubRasterizer(render_context, 0)
dataset = dataset_cls(cfg, zarr_dataset, rasterizer, None)
indices = dataset.get_scene_indices(scene_idx)
subdata = Subset(dataset, indices)
for _ in subdata: # type: ignore
pass
def test_compute_mse_error(tmp_path: Path, zarr_dataset: ChunkedDataset, cfg: dict) -> None:
render_context = RenderContext(
np.asarray((10, 10)),
np.asarray((0.25, 0.25)),
np.asarray((0.5, 0.5)),
set_origin_to_bottom=cfg["raster_params"]["set_origin_to_bottom"],
)
rast = StubRasterizer(render_context)
dataset = AgentDataset(cfg, zarr_dataset, rast)
gt_coords = []
gt_avails = []
timestamps = []
track_ids = []
for idx, el in enumerate(dataset): # type: ignore
gt_coords.append(el["target_positions"])
gt_avails.append(el["target_availabilities"])
timestamps.append(el["timestamp"])
track_ids.append(el["track_id"])
if idx == 100:
break # speed up test
gt_coords = np.asarray(gt_coords)
gt_avails = np.asarray(gt_avails)
timestamps = np.asarray(timestamps)
track_ids = np.asarray(track_ids)
# test same values error
write_gt_csv(str(tmp_path / "gt1.csv"), timestamps, track_ids, gt_coords, gt_avails)
write_pred_csv(str(tmp_path / "pred1.csv"), timestamps, track_ids, gt_coords, confs=None)
metrics = compute_metrics_csv(str(tmp_path / "gt1.csv"), str(tmp_path / "pred1.csv"), [neg_multi_log_likelihood])
for metric_value in metrics.values():
assert np.all(metric_value == 0.0)
# test different values error
pred_coords = gt_coords.copy()
pred_coords += np.random.randn(*pred_coords.shape)
write_pred_csv(str(tmp_path / "pred3.csv"), timestamps, track_ids, pred_coords, confs=None)
metrics = compute_metrics_csv(str(tmp_path / "gt1.csv"), str(tmp_path / "pred3.csv"), [neg_multi_log_likelihood])
for metric_value in metrics.values():
assert np.any(metric_value > 0.0)
# test invalid conf by removing lines in gt1
with open(str(tmp_path / "pred4.csv"), "w") as fp:
lines = open(str(tmp_path / "pred1.csv")).readlines()
fp.writelines(lines[:-10])
with pytest.raises(ValueError):
compute_metrics_csv(str(tmp_path / "gt1.csv"), str(tmp_path / "pred4.csv"), [neg_multi_log_likelihood])
def test_frame_index_interval(dataset_cls: Callable, frame_idx: int,
zarr_dataset: ChunkedDataset, cfg: dict) -> None:
render_context = RenderContext(
np.asarray((100, 100)),
np.asarray((0.25, 0.25)),
np.asarray((0.5, 0.5)),
set_origin_to_bottom=cfg["raster_params"]["set_origin_to_bottom"],
)
rasterizer = StubRasterizer(render_context)
dataset = dataset_cls(cfg, zarr_dataset, rasterizer, None)
indices = dataset.get_frame_indices(frame_idx)
subdata = Subset(dataset, indices)
for _ in subdata: # type: ignore
pass
def build_custom_rasterizer(cfg: dict, data_manager):
raster_cfg = cfg["raster_params"]
map_type = raster_cfg["map_type"]
if map_type not in ("semantic_fast", "box_fast", "box_semantic_fast"):
return build_rasterizer(cfg, data_manager)
raster_cfg = cfg["raster_params"]
map_type = raster_cfg["map_type"]
dataset_meta_key = raster_cfg["dataset_meta_key"]
render_context = RenderContext(
raster_size_px=np.array(raster_cfg["raster_size"]),
pixel_size_m=np.array(raster_cfg["pixel_size"]),
center_in_raster_ratio=np.array(raster_cfg["ego_center"]),
)
filter_agents_threshold = raster_cfg["filter_agents_threshold"]
history_num_frames = cfg["model_params"]["history_num_frames"]
history_box_frames = cfg["model_params"].get("history_box_frames", [])
semantic_map_filepath = data_manager.require(raster_cfg["semantic_map_key"])
try:
dataset_meta = _load_metadata(dataset_meta_key, data_manager)
world_to_ecef = np.array(dataset_meta["world_to_ecef"], dtype=np.float64)
except (
KeyError,
FileNotFoundError,
): # TODO remove when new dataset version is available
world_to_ecef = get_hardcoded_world_to_ecef()
if map_type == "semantic_fast":
return FastSemanticRasterizer(render_context, semantic_map_filepath, world_to_ecef)
elif map_type == "box_fast":
return FastBoxRasterizer(
render_context,
filter_agents_threshold,
history_num_frames,
history_box_frames,
)
elif map_type == "box_semantic_fast":
return FastSemBoxRasterizer(
render_context=render_context,
filter_agents_threshold=filter_agents_threshold,
world_to_ecef=world_to_ecef,
semantic_map_path=semantic_map_filepath,
history_num_frames=history_num_frames,
box_history_frames=history_box_frames,
)
def from_cfg(cfg, data_manager=None):
raster_cfg = cfg["raster_params"]
# map_type = raster_cfg["map_type"]
# dataset_meta_key = raster_cfg["dataset_meta_key"]
render_context = RenderContext(
raster_size_px=np.array(raster_cfg["raster_size"]),
pixel_size_m=np.array(raster_cfg["pixel_size"]),
center_in_raster_ratio=np.array(raster_cfg["ego_center"]),
)
# raster_size: Tuple[int, int] = cast(Tuple[int, int], tuple(raster_cfg["raster_size"]))
# pixel_size = np.array(raster_cfg["pixel_size"])
# ego_center = np.array(raster_cfg["ego_center"])
filter_agents_threshold = raster_cfg["filter_agents_threshold"]
history_num_frames = cfg["model_params"]["history_num_frames"]
return VelocityBoxRasterizer(render_context, filter_agents_threshold, history_num_frames)
def create_generate_agent_sample_fixing_yaw_partial(cfg, rasterizer):
render_context = RenderContext(
raster_size_px=np.array(cfg["raster_params"]["raster_size"]),
pixel_size_m=np.array(cfg["raster_params"]["pixel_size"]),
center_in_raster_ratio=np.array(cfg["raster_params"]["ego_center"]),
)
# SemanticRasterizer の一部機能を作るため build する
# TODO(kaizaburo): LocalDataManager を外から受け取るか、そもそも SemanticRasterizer を使わない
assert "L5KIT_DATA_FOLDER" in os.environ
sat_rast_cfg = copy.deepcopy(cfg)
sat_rast_cfg["raster_params"]["map_type"] = "semantic_debug"
sat_rast: SemanticRasterizer = build_rasterizer(sat_rast_cfg,
LocalDataManager(None))
# Tuned by optuna
default_yaw_fix_params = {
"angle_diff_th_in_degree": 30.00639299068446,
"distance_th_in_agent": 0.031291713326932744,
"num_use_history": 4,
"lane_polygon_buffer": 0.8153613050209113
}
yaw_fix_params = cfg["raster_params"].get("yaw_fix_params",
default_yaw_fix_params)
sample_function = partial(
generate_agent_sample_fixing_yaw,
render_context=render_context,
history_num_frames=cfg["model_params"]["history_num_frames"],
history_step_size=cfg["model_params"]["history_step_size"],
future_num_frames=cfg["model_params"]["future_num_frames"],
future_step_size=cfg["model_params"]["future_step_size"],
filter_agents_threshold=cfg["raster_params"]
["filter_agents_threshold"],
sat_rast=sat_rast,
rasterizer=rasterizer,
perturbation=None,
angle_diff_th_in_degree=yaw_fix_params["angle_diff_th_in_degree"],
distance_th_in_agent=yaw_fix_params["distance_th_in_agent"],
num_use_history=yaw_fix_params["num_use_history"],
lane_polygon_buffer=yaw_fix_params["lane_polygon_buffer"],
)
return sample_function
def test_coordinates_straight_road(zarr_dataset: ChunkedDataset,
cfg: dict) -> None:
# on a straight road `target_positions` should increase on x only
render_context = RenderContext(
np.asarray(cfg["raster_params"]["raster_size"]),
np.asarray(cfg["raster_params"]["pixel_size"]),
np.asarray(cfg["raster_params"]["ego_center"]),
)
dataset = EgoDataset(
cfg,
zarr_dataset,
StubRasterizer(
render_context,
0.5,
),
)
# get first prediction and first 50 centroids
centroids = []
preds = []
preds_world = []
for idx in range(50):
data = dataset[idx]
if idx == 0:
preds = data["target_positions"]
preds_world = transform_points(
preds, np.linalg.inv(data["agent_from_world"]))
centroids.append(data["centroid"][:2])
centroids = np.stack(centroids)
# compute XY variances for preds and centroids
var_preds = np.var(preds, 0, ddof=1)
var_centroids = np.var(centroids, 0, ddof=1)
assert var_preds[1] / var_preds[
0] < 0.001 # variance on Y is way lower than on X
assert var_centroids[1] / var_centroids[
0] > 0.9 # variance on Y is similar to X
# check similarity between coordinates
assert np.allclose(preds_world[:-1], centroids[1:])
def create_generate_agent_sample_tl_history_partial(cfg, rasterizer):
render_context = RenderContext(
raster_size_px=np.array(cfg["raster_params"]["raster_size"]),
pixel_size_m=np.array(cfg["raster_params"]["pixel_size"]),
center_in_raster_ratio=np.array(cfg["raster_params"]["ego_center"]),
)
sample_function = partial(
generate_agent_sample_tl_history,
render_context=render_context,
history_num_frames=cfg["model_params"]["history_num_frames"],
history_step_size=cfg["model_params"]["history_step_size"],
future_num_frames=cfg["model_params"]["future_num_frames"],
future_step_size=cfg["model_params"]["future_step_size"],
filter_agents_threshold=cfg["raster_params"]
["filter_agents_threshold"],
rasterizer=rasterizer,
perturbation=None,
)
return sample_function
def export_zarr_to_csv(
zarr_dataset: ChunkedDataset,
csv_file_path: str,
future_num_frames: int,
filter_agents_threshold: float,
step_time: float = 0.1,
agents_mask: Optional[np.array] = None,
) -> None:
"""Produces a csv file containing the ground truth from a zarr file.
Arguments:
zarr_dataset (np.ndarray): The open zarr dataset.
csv_file_path (str): File path to write a CSV to.
future_num_frames (int): Amount of future displacements we want.
filter_agents_threshold (float): Value between 0 and 1 to use as cutoff value for agent filtering
agents_mask (Optional[np.array]): a boolean mask of shape (len(zarr_dataset.agents)) which will be used
instead of computing the agents_mask
"""
cfg = {
"raster_params": {
"pixel_size": np.asarray((0.25, 0.25)),
"raster_size": (100, 100),
"filter_agents_threshold": filter_agents_threshold,
"disable_traffic_light_faces": True,
"ego_center": np.asarray((0.5, 0.5)),
"set_origin_to_bottom": True,
},
"model_params": {
"history_num_frames": 0,
"future_num_frames": future_num_frames,
"step_time": step_time
},
}
render_context = RenderContext(
np.asarray(cfg["raster_params"]["raster_size"]),
cfg["raster_params"]["pixel_size"],
cfg["raster_params"]["ego_center"],
cfg["raster_params"]["set_origin_to_bottom"],
)
rasterizer = StubRasterizer(render_context)
dataset = AgentDataset(cfg=cfg,
zarr_dataset=zarr_dataset,
rasterizer=rasterizer,
agents_mask=agents_mask)
future_coords_offsets = []
target_availabilities = []
timestamps = []
agent_ids = []
for el in tqdm(dataset, desc="extracting GT"): # type: ignore
# convert agent coordinates to world offsets
offsets = transform_points(el["target_positions"],
el["world_from_agent"]) - el["centroid"][:2]
future_coords_offsets.append(offsets)
timestamps.append(el["timestamp"])
agent_ids.append(el["track_id"])
target_availabilities.append(el["target_availabilities"])
write_gt_csv(
csv_file_path,
np.asarray(timestamps),
np.asarray(agent_ids),
np.asarray(future_coords_offsets),
np.asarray(target_availabilities),
)
def generate_agent_sample_fixing_yaw(
state_index: int,
frames: np.ndarray,
agents: np.ndarray,
tl_faces: np.ndarray,
selected_track_id: Optional[int],
render_context: RenderContext,
history_num_frames: int,
history_step_size: int,
future_num_frames: int,
future_step_size: int,
filter_agents_threshold: float,
sat_rast: SemanticRasterizer,
rasterizer: Optional[Rasterizer] = None,
perturbation: Optional[Perturbation] = None,
angle_diff_th_in_degree: float = 60,
distance_th_in_agent: float = 0.5,
num_use_history: int = 2,
lane_polygon_buffer: float = 1.5,
) -> dict:
"""Generates the inputs and targets to train a deep prediction model. A deep prediction model takes as input
the state of the world (here: an image we will call the "raster"), and outputs where that agent will be some
seconds into the future.
This function has a lot of arguments and is intended for internal use, you should try to use higher level classes
and partials that use this function.
Args:
state_index (int): The anchor frame index, i.e. the "current" timestep in the scene
frames (np.ndarray): The scene frames array, can be numpy array or a zarr array
agents (np.ndarray): The full agents array, can be numpy array or a zarr array
tl_faces (np.ndarray): The full traffic light faces array, can be numpy array or a zarr array
selected_track_id (Optional[int]): Either None for AV, or the ID of an agent that you want to
predict the future of. This agent is centered in the raster and the returned targets are derived from
their future states.
raster_size (Tuple[int, int]): Desired output raster dimensions
pixel_size (np.ndarray): Size of one pixel in the real world
ego_center (np.ndarray): Where in the raster to draw the ego, [0.5,0.5] would be the center
history_num_frames (int): Amount of history frames to draw into the rasters
history_step_size (int): Steps to take between frames, can be used to subsample history frames
future_num_frames (int): Amount of history frames to draw into the rasters
future_step_size (int): Steps to take between targets into the future
filter_agents_threshold (float): Value between 0 and 1 to use as cutoff value for agent filtering
based on their probability of being a relevant agent
rasterizer (Optional[Rasterizer]): Rasterizer of some sort that draws a map image
perturbation (Optional[Perturbation]): Object that perturbs the input and targets, used
to train models that can recover from slight divergence from training set data
Raises:
ValueError: A ValueError is returned if the specified ``selected_track_id`` is not present in the scene
or was filtered by applying the ``filter_agent_threshold`` probability filtering.
Returns:
dict: a dict object with the raster array, the future offset coordinates (meters),
the future yaw angular offset, the future_availability as a binary mask
"""
# the history slice is ordered starting from the latest frame and goes backward in time., ex. slice(100, 91, -2)
history_slice = get_history_slice(state_index,
history_num_frames,
history_step_size,
include_current_state=True)
future_slice = get_future_slice(state_index, future_num_frames,
future_step_size)
history_frames = frames[history_slice].copy(
) # copy() required if the object is a np.ndarray
future_frames = frames[future_slice].copy()
sorted_frames = np.concatenate(
(history_frames[::-1], future_frames)) # from past to future
# get agents (past and future)
agent_slice = get_agents_slice_from_frames(sorted_frames[0],
sorted_frames[-1])
agents = agents[agent_slice].copy(
) # this is the minimum slice of agents we need
history_frames[
"agent_index_interval"] -= agent_slice.start # sync interval with the agents array
future_frames[
"agent_index_interval"] -= agent_slice.start # sync interval with the agents array
history_agents = filter_agents_by_frames(history_frames, agents)
future_agents = filter_agents_by_frames(future_frames, agents)
tl_slice = get_tl_faces_slice_from_frames(
history_frames[-1], history_frames[0]) # -1 is the farthest
# sync interval with the traffic light faces array
history_frames["traffic_light_faces_index_interval"] -= tl_slice.start
history_tl_faces = filter_tl_faces_by_frames(history_frames,
tl_faces[tl_slice].copy())
if perturbation is not None:
history_frames, future_frames = perturbation.perturb(
history_frames=history_frames, future_frames=future_frames)
# State you want to predict the future of.
cur_frame = history_frames[0]
cur_agents = history_agents[0]
if selected_track_id is None:
agent_centroid_m = cur_frame["ego_translation"][:2]
agent_yaw_rad = rotation33_as_yaw(cur_frame["ego_rotation"])
agent_extent_m = np.asarray(
(EGO_EXTENT_LENGTH, EGO_EXTENT_WIDTH, EGO_EXTENT_HEIGHT))
selected_agent = None
yaw_fixed = False
else:
# this will raise IndexError if the agent is not in the frame or under agent-threshold
# this is a strict error, we cannot recover from this situation
try:
agent = filter_agents_by_track_id(
filter_agents_by_labels(cur_agents, filter_agents_threshold),
selected_track_id)[0]
except IndexError:
raise ValueError(
f" track_id {selected_track_id} not in frame or below threshold"
)
# Try to fix agent's yaw
agent, yaw_fixed = _try_to_fix_agent_yaw(
sat_rast=sat_rast,
agent=agent,
selected_track_id=selected_track_id,
history_frames=history_frames,
history_agents=history_agents,
history_num_frames=history_num_frames,
angle_diff_th_in_degree=angle_diff_th_in_degree,
distance_th_in_agent=distance_th_in_agent,
num_use_history=num_use_history,
lane_polygon_buffer=lane_polygon_buffer,
)
agent_centroid_m = agent["centroid"]
agent_yaw_rad = float(agent["yaw"])
agent_extent_m = agent["extent"]
selected_agent = agent
input_im = (None if not rasterizer else rasterizer.rasterize(
history_frames, history_agents, history_tl_faces, selected_agent))
world_from_agent = compute_agent_pose(agent_centroid_m, agent_yaw_rad)
agent_from_world = np.linalg.inv(world_from_agent)
raster_from_world = render_context.raster_from_world(
agent_centroid_m, agent_yaw_rad)
future_coords_offset, future_yaws_offset, future_availability = _create_targets_for_deep_prediction(
future_num_frames, future_frames, selected_track_id, future_agents,
agent_from_world, agent_yaw_rad)
# history_num_frames + 1 because it also includes the current frame
history_coords_offset, history_yaws_offset, history_availability = _create_targets_for_deep_prediction(
history_num_frames + 1, history_frames, selected_track_id,
history_agents, agent_from_world, agent_yaw_rad)
return {
"image": input_im,
"target_positions": future_coords_offset,
"target_yaws": future_yaws_offset,
"target_availabilities": future_availability,
"history_positions": history_coords_offset,
"history_yaws": history_yaws_offset,
"history_availabilities": history_availability,
"world_to_image": raster_from_world, # TODO deprecate
"raster_from_agent": raster_from_world @ world_from_agent,
"raster_from_world": raster_from_world,
"agent_from_world": agent_from_world,
"world_from_agent": world_from_agent,
"centroid": agent_centroid_m,
"yaw": agent_yaw_rad,
"extent": agent_extent_m,
"yaw_fixed": yaw_fixed,
}
def generate_multi_agent_sample(
state_index: int,
frames: np.ndarray,
agents: np.ndarray,
tl_faces: np.ndarray,
selected_track_id: Optional[int],
render_context: RenderContext,
history_num_frames: int,
history_step_size: int,
future_num_frames: int,
future_step_size: int,
filter_agents_threshold: float,
rasterizer: Optional[Rasterizer] = None,
perturbation: Optional[Perturbation] = None,
min_frame_history: int = MIN_FRAME_HISTORY,
min_frame_future: int = MIN_FRAME_FUTURE,
) -> dict:
"""Generates the inputs and targets to train a deep prediction model. A deep prediction model takes as input
the state of the world (here: an image we will call the "raster"), and outputs where that agent will be some
seconds into the future.
This function has a lot of arguments and is intended for internal use, you should try to use higher level classes
and partials that use this function.
Args:
state_index (int): The anchor frame index, i.e. the "current" timestep in the scene
frames (np.ndarray): The scene frames array, can be numpy array or a zarr array
agents (np.ndarray): The full agents array, can be numpy array or a zarr array
tl_faces (np.ndarray): The full traffic light faces array, can be numpy array or a zarr array
selected_track_id (Optional[int]): Either None for AV, or the ID of an agent that you want to
predict the future of. This agent is centered in the raster and the returned targets are derived from
their future states.
render_context (RenderContext):
raster_size (Tuple[int, int]): Desired output raster dimensions
pixel_size (np.ndarray): Size of one pixel in the real world
ego_center (np.ndarray): Where in the raster to draw the ego, [0.5,0.5] would be the center
history_num_frames (int): Amount of history frames to draw into the rasters
history_step_size (int): Steps to take between frames, can be used to subsample history frames
future_num_frames (int): Amount of history frames to draw into the rasters
future_step_size (int): Steps to take between targets into the future
filter_agents_threshold (float): Value between 0 and 1 to use as cutoff value for agent filtering
based on their probability of being a relevant agent
rasterizer (Optional[Rasterizer]): Rasterizer of some sort that draws a map image
perturbation (Optional[Perturbation]): Object that perturbs the input and targets, used
to train models that can recover from slight divergence from training set data
Raises:
ValueError: A ValueError is returned if the specified ``selected_track_id`` is not present in the scene
or was filtered by applying the ``filter_agent_threshold`` probability filtering.
Returns:
dict: a dict object with the raster array, the future offset coordinates (meters),
the future yaw angular offset, the future_availability as a binary mask
"""
# the history slice is ordered starting from the latest frame and goes backward in time., ex. slice(100, 91, -2)
history_slice = get_history_slice(state_index,
history_num_frames,
history_step_size,
include_current_state=True)
future_slice = get_future_slice(state_index, future_num_frames,
future_step_size)
history_frames = frames[history_slice].copy(
) # copy() required if the object is a np.ndarray
future_frames = frames[future_slice].copy()
sorted_frames = np.concatenate(
(history_frames[::-1], future_frames)) # from past to future
# get agents (past and future)
agent_slice = get_agents_slice_from_frames(sorted_frames[0],
sorted_frames[-1])
agents = agents[agent_slice].copy(
) # this is the minimum slice of agents we need
history_frames[
"agent_index_interval"] -= agent_slice.start # sync interval with the agents array
future_frames[
"agent_index_interval"] -= agent_slice.start # sync interval with the agents array
history_agents = filter_agents_by_frames(history_frames, agents)
future_agents = filter_agents_by_frames(future_frames, agents)
try:
tl_slice = get_tl_faces_slice_from_frames(
history_frames[-1], history_frames[0]) # -1 is the farthest
# sync interval with the traffic light faces array
history_frames["traffic_light_faces_index_interval"] -= tl_slice.start
history_tl_faces = filter_tl_faces_by_frames(history_frames,
tl_faces[tl_slice].copy())
except ValueError:
history_tl_faces = [
np.empty(0, dtype=TL_FACE_DTYPE) for _ in history_frames
]
if perturbation is not None:
history_frames, future_frames = perturbation.perturb(
history_frames=history_frames, future_frames=future_frames)
# State you want to predict the future of.
cur_frame = history_frames[0]
cur_agents = history_agents[0]
cur_agents = filter_agents_by_labels(cur_agents, filter_agents_threshold)
agent_track_ids_u64 = cur_agents["track_id"]
# uint64 --> int64
agent_track_ids = agent_track_ids_u64.astype(np.int64)
assert np.alltrue(agent_track_ids == agent_track_ids_u64)
agent_track_ids = np.concatenate(
[np.array([-1], dtype=np.int64), agent_track_ids])
# Draw image with Ego car in center
selected_agent = None
input_im = (None if not rasterizer else rasterizer.rasterize(
history_frames, history_agents, history_tl_faces, selected_agent))
future_coords_offset_list = []
future_yaws_offset_list = []
future_availability_list = []
history_coords_offset_list = []
history_yaws_offset_list = []
history_availability_list = []
agent_centroid_list = []
agent_yaw_list = []
agent_extent_list = []
filtered_track_ids_list = []
for selected_track_id in agent_track_ids:
if selected_track_id == -1:
agent_centroid = cur_frame["ego_translation"][:2]
agent_yaw_rad = rotation33_as_yaw(cur_frame["ego_rotation"])
agent_extent = np.asarray(
(EGO_EXTENT_LENGTH, EGO_EXTENT_WIDTH, EGO_EXTENT_HEIGHT))
world_from_agent = compute_agent_pose(agent_centroid,
agent_yaw_rad)
agent_from_world = np.linalg.inv(world_from_agent)
raster_from_world = render_context.raster_from_world(
agent_centroid, agent_yaw_rad)
agent_origin = np.zeros((2, ), dtype=np.float32)
else:
# this will raise IndexError if the agent is not in the frame or under agent-threshold
# this is a strict error, we cannot recover from this situation
try:
agent = filter_agents_by_track_id(cur_agents,
selected_track_id)[0]
except IndexError:
raise ValueError(
f" track_id {selected_track_id} not in frame or below threshold"
)
agent_centroid = agent["centroid"]
agent_yaw_rad = agent["yaw"]
agent_extent = agent["extent"]
agent_origin = transform_point(agent_centroid, agent_from_world)
future_coords_offset, future_yaws_offset, future_availability = _create_targets_for_deep_prediction(
future_num_frames, future_frames, selected_track_id, future_agents,
agent_from_world, agent_yaw_rad, agent_origin)
if selected_track_id != -1 and np.sum(
future_availability) < min_frame_future:
# Not enough future to predict, skip this agent.
continue
# history_num_frames + 1 because it also includes the current frame
history_coords_offset, history_yaws_offset, history_availability = _create_targets_for_deep_prediction(
history_num_frames + 1, history_frames, selected_track_id,
history_agents, agent_from_world, agent_yaw_rad, agent_origin)
if selected_track_id != -1 and np.sum(
history_availability) < min_frame_history:
# Not enough history to predict, skip this agent.
continue
future_coords_offset_list.append(future_coords_offset)
future_yaws_offset_list.append(future_yaws_offset)
future_availability_list.append(future_availability)
history_coords_offset_list.append(history_coords_offset)
history_yaws_offset_list.append(history_yaws_offset)
history_availability_list.append(history_availability)
agent_centroid_list.append(agent_centroid)
agent_yaw_list.append(agent_yaw_rad)
agent_extent_list.append(agent_extent)
filtered_track_ids_list.append(selected_track_id)
# Get pixel coordinate
agent_centroid_array = np.array(agent_centroid_list)
agent_centroid_in_pixel = transform_points(agent_centroid_array,
raster_from_world)
return {
"image": input_im, # (h, w, ch)
# --- All below is in world coordinate ---
"target_positions":
np.array(future_coords_offset_list), # (n_agents, num_frames, 2)
"target_yaws":
np.array(future_yaws_offset_list), # (n_agents, num_frames, 1)
"target_availabilities":
np.array(future_availability_list), # (n_agents, num_frames)
"history_positions":
np.array(history_coords_offset_list), # (n_agents, num_frames, 2)
"history_yaws":
np.array(history_yaws_offset_list), # (n_agents, num_frames, 1)
"history_availabilities":
np.array(history_availability_list), # (n_agents, num_frames)
# "world_to_image": raster_from_world, # (3, 3)
"raster_from_world": raster_from_world, # (3, 3)
"centroid": agent_centroid_array, # (n_agents, 2)
"yaw": np.array(agent_yaw_list), # (n_agents, 1)
"extent": np.array(agent_extent_list), # (n_agents, 3)
"track_ids": np.array(filtered_track_ids_list), # (n_agents)
"centroid_pixel": agent_centroid_in_pixel, # (n_agents, 2)
}