def get_box_nodes(self, history_frames: np.ndarray,
history_agents: List[np.ndarray],
agent: Optional[np.ndarray],
raster_from_world: np.ndarray) -> Dict[str, List]:
"""
remove node not in the latest
"""
all_agents_info = []
ego_info = []
# loop over historical frames
for i, (frame, agents) in enumerate(zip(history_frames,
history_agents)):
agents = filter_agents_by_labels(agents,
self.filter_agents_threshold)
av_agent = get_ego_as_agent(frame).astype(agents.dtype)
agents = np.concatenate([agents, av_agent])
agent_ego = filter_agents_by_track_id(agents, agent["track_id"])
agents = self.__class__.remove_agents_by_track_id(
agents, track_id=agent["track_id"])
agent_ego["centroid"] = transform_points(agent_ego["centroid"],
raster_from_world)
agents["centroid"] = transform_points(agents["centroid"],
raster_from_world)
all_agents_info.append(agents)
ego_info.append(agent_ego)
return {
self.__class__.AGENTS: all_agents_info,
self.__class__.EGO: ego_info
}
def rasterize(
self,
history_frames: np.ndarray,
history_agents: List[np.ndarray],
history_tl_faces: List[np.ndarray],
agent: Optional[np.ndarray] = None,
) -> np.ndarray:
# all frames are drawn relative to this one"
frame = history_frames[0]
if agent is None:
ego_translation_m = frame["ego_translation"]
ego_yaw_rad = rotation33_as_yaw(frame["ego_rotation"])
else:
ego_translation_m = np.append(agent["centroid"], history_frames[0]["ego_translation"][-1])
ego_yaw_rad = agent["yaw"]
if self.pixel_size[0] != self.pixel_size[1]:
raise NotImplementedError("No support for non squared pixels yet")
raster_from_world = self.render_context.raster_from_world(ego_translation_m, ego_yaw_rad)
# this ensures we always end up with fixed size arrays, +1 is because current time is also in the history
out_shape = (self.raster_size[1], self.raster_size[0], self.history_num_frames + 1)
agents_images = np.zeros(out_shape, dtype=np.uint8)
ego_images = np.zeros(out_shape, dtype=np.uint8)
velocity_images = np.zeros(out_shape, dtype=np.uint8)
for i, (frame, agents) in enumerate(zip(history_frames, history_agents)):
agents = filter_agents_by_labels(agents, self.filter_agents_threshold)
# note the cast is for legacy support of dataset before April 2020
av_agent = get_ego_as_agent(frame).astype(agents.dtype)
velocity_image = draw_velocity(self.raster_size, raster_from_world, np.append(agents, av_agent), 255)
if agent is None:
agents_image = draw_boxes(self.raster_size, raster_from_world, agents, 255)
ego_image = draw_boxes(self.raster_size, raster_from_world, av_agent, 255)
else:
agent_ego = filter_agents_by_track_id(agents, agent["track_id"])
if len(agent_ego) == 0: # agent not in this history frame
agents_image = draw_boxes(self.raster_size, raster_from_world, np.append(agents, av_agent), 255)
ego_image = np.zeros_like(agents_image)
else: # add av to agents and remove the agent from agents
agents = agents[agents != agent_ego[0]]
agents_image = draw_boxes(self.raster_size, raster_from_world, np.append(agents, av_agent), 255)
ego_image = draw_boxes(self.raster_size, raster_from_world, agent_ego, 255)
agents_images[..., i] = agents_image
ego_images[..., i] = ego_image
velocity_images[..., i] = velocity_image
# combine such that the image consists of [agent_t, agent_t-1, agent_t-2, ego_t, ego_t-1, ego_t-2]
out_im = np.concatenate((agents_images, ego_images, velocity_images), -1)
assert out_im.shape[-1] == self.raster_channels
return out_im.astype(np.float32) / 255
def rasterize_box(
self,
history_frames: np.ndarray,
history_agents: List[np.ndarray],
history_tl_faces: List[np.ndarray],
agent: Optional[np.ndarray] = None,
svg=False, svg_args=None,
) -> th.Union[dict]:
# all frames are drawn relative to this one"
frame = history_frames[0]
if agent is None:
ego_translation_m = history_frames[0]["ego_translation"]
ego_yaw_rad = rotation33_as_yaw(frame["ego_rotation"])
else:
ego_translation_m = np.append(agent["centroid"], history_frames[0]["ego_translation"][-1])
ego_yaw_rad = agent["yaw"]
svg_args = svg_args or dict()
raster_from_world = self.render_context.raster_from_world(ego_translation_m, ego_yaw_rad)
raster_size = self.render_context.raster_size_px
# this ensures we always end up with fixed size arrays, +1 is because current time is also in the history
res = dict(ego=list(), agents=defaultdict(list))
for i, (frame, agents) in enumerate(zip(history_frames, history_agents)):
# print('history index', i)
agents = filter_agents_by_labels(agents, self.filter_agents_threshold)
# note the cast is for legacy support of dataset before April 2020
av_agent = get_ego_as_agent(frame).astype(agents.dtype)
if agent is None:
add_agents(res['agents'], av_agent)
res['ego'].append(av_agent[0]["centroid"][:2])
else:
agent_ego = filter_agents_by_track_id(agents, agent["track_id"])
if len(agent_ego) == 0: # agent not in this history frame
add_agents(res['agents'], np.append(agents, av_agent))
else: # add av to agents and remove the agent from agents
agents = agents[agents != agent_ego[0]]
add_agents(res['agents'], np.append(agents, av_agent))
res['ego'].append(agent_ego[0]["centroid"][:2])
tolerance = svg_args.get('tolerance', 20.)
_ego = normalize_line(
cv2_subpixel(transform_points(np.array(res['ego']).reshape((-1, 2)), raster_from_world)))
res['ego'] = crop_tensor(_ego, raster_size)
ego_grad = calc_max_grad(res['ego'])
res['agents'] = [normalize_line(cv2_subpixel(transform_points(np.array(path).reshape((-1, 2)), raster_from_world))
) for idx, path in res['agents'].items()]
res['agents'] = [
crop_tensor(path, raster_size) for path in res['agents'] if not is_noisy(path, ego_grad, tolerance)]
res['agents'] = [path for path in res['agents'] if len(path)]
if svg:
res['path'] = torch.cat([linear_path_to_tensor(path, svg_args.get('pad_val', -1)) for path in res['agents']
] + [linear_path_to_tensor(res['ego'], svg_args.get('pad_val', -1))], 0)
res['path_type'] = [path_type_to_number('agent')] * len(res['agents']) + [path_type_to_number('ego')]
return res
def rasterize(
self,
history_frames: np.ndarray,
history_agents: List[np.ndarray],
history_tl_faces: List[np.ndarray],
agent: Optional[np.ndarray] = None,
) -> np.ndarray:
# all frames are drawn relative to this one"
frame = history_frames[0]
if agent is None:
ego_translation_m = history_frames[0]["ego_translation"]
ego_yaw_rad = rotation33_as_yaw(frame["ego_rotation"])
else:
ego_translation_m = np.append(agent["centroid"], history_frames[0]["ego_translation"][-1])
ego_yaw_rad = agent["yaw"]
raster_from_world = self.render_context.raster_from_world(ego_translation_m, ego_yaw_rad)
# this ensures we always end up with fixed size arrays, +1 is because current time is also in the history
out_shape = (self.raster_size[1], self.raster_size[0], self.history_num_frames + 1)
agents_unknown_images = np.zeros(out_shape, dtype=np.uint8)
agents_car_images = np.zeros(out_shape, dtype=np.uint8)
agents_cyclist_images = np.zeros(out_shape, dtype=np.uint8)
agents_pedestrian_images = np.zeros(out_shape, dtype=np.uint8)
ego_images = np.zeros(out_shape, dtype=np.uint8)
if self.enable_selected_agent_channels:
assert agent is not None
selected_agent_images = np.zeros(out_shape, dtype=np.uint8)
for i, (frame, agents) in enumerate(zip(history_frames, history_agents)):
# TODO: filter_agents_threshold is not used.
# Ignore filter_agents_threshold now
threshold = 0.5
agents_unknown = agents[agents["label_probabilities"][:, UNKNOWN_LABEL_INDEX] > threshold]
agents_car = agents[agents["label_probabilities"][:, CAR_LABEL_INDEX] > threshold]
agents_cyclist = agents[agents["label_probabilities"][:, CYCLIST_LABEL_INDEX] > threshold]
agents_pedestrian = agents[agents["label_probabilities"][:, PEDESTRIAN_LABEL_INDEX] > threshold]
# agents = filter_agents_by_labels(agents, self.filter_agents_threshold)
# note the cast is for legacy support of dataset before April 2020
agents_ego = get_ego_as_agent(frame).astype(agents.dtype)
agents_unknown_image = draw_boxes(self.raster_size, raster_from_world, agents_unknown, 255)
agents_car_image = draw_boxes(self.raster_size, raster_from_world, agents_car, 255)
agents_cyclist_image = draw_boxes(self.raster_size, raster_from_world, agents_cyclist, 255)
agents_pedestrian_image = draw_boxes(self.raster_size, raster_from_world, agents_pedestrian, 255)
ego_image = draw_boxes(self.raster_size, raster_from_world, agents_ego, 255)
agents_unknown_images[..., i] = agents_unknown_image
agents_car_images[..., i] = agents_car_image
agents_cyclist_images[..., i] = agents_cyclist_image
agents_pedestrian_images[..., i] = agents_pedestrian_image
ego_images[..., i] = ego_image
if self.enable_selected_agent_channels:
assert agent is not None
selected_agent = filter_agents_by_track_id(agents, agent["track_id"])
if len(selected_agent) == 0: # agent not in this history frame
selected_agent_image = np.zeros_like(ego_image)
else: # add av to agents and remove the agent from agents
selected_agent_image = draw_boxes(self.raster_size, raster_from_world, selected_agent, 255)
selected_agent_images[..., i] = selected_agent_image
images = [
agents_unknown_images,
agents_car_images,
agents_cyclist_images,
agents_pedestrian_images,
ego_images
]
if self.enable_selected_agent_channels:
images.append(selected_agent_images)
out_im = np.concatenate(images, axis=-1)
assert out_im.shape[-1] == self.raster_channels
return out_im.astype(np.float32) / 255
def rasterize(
self,
history_frames: np.ndarray,
history_agents: List[np.ndarray],
history_tl_faces: List[np.ndarray],
agent: Optional[np.ndarray] = None,
) -> np.ndarray:
# all frames are drawn relative to this one"
frame = history_frames[0]
if agent is None:
ego_translation_m = history_frames[0]["ego_translation"]
ego_yaw_rad = rotation33_as_yaw(frame["ego_rotation"])
else:
ego_translation_m = np.append(agent["centroid"], history_frames[0]["ego_translation"][-1])
ego_yaw_rad = agent["yaw"]
raster_from_world = self.render_context.raster_from_world(ego_translation_m, ego_yaw_rad)
# this ensures we always end up with fixed size arrays, +1 is because current time is also in the history
# TODO: Change shape
# out_shape = (self.raster_size[1], self.raster_size[0], self.history_num_frames + 1)
out_shape = (self.raster_size[1], self.raster_size[0], 1)
# agents_images = np.zeros(out_shape, dtype=np.uint8)
# ego_images = np.zeros(out_shape, dtype=np.uint8)
agents_images = np.zeros(out_shape, dtype=np.float32)
ego_images = np.zeros(out_shape, dtype=np.float32)
n = history_frames.size
weights = (np.logspace(0, 1, n, base=15) / n)[::-1][:, None, None]
for i, (frame, agents) in enumerate(zip(history_frames, history_agents)):
agents = filter_agents_by_labels(agents, self.filter_agents_threshold)
# note the cast is for legacy support of dataset before April 2020
av_agent = get_ego_as_agent(frame).astype(agents.dtype)
if agent is None:
agents_image = draw_boxes(self.raster_size, raster_from_world, agents, 255)
ego_image = draw_boxes(self.raster_size, raster_from_world, av_agent, 255)
else:
agent_ego = filter_agents_by_track_id(agents, agent["track_id"])
if len(agent_ego) == 0: # agent not in this history frame
agents_image = draw_boxes(self.raster_size, raster_from_world, np.append(agents, av_agent), 255)
ego_image = np.zeros_like(agents_image)
else: # add av to agents and remove the agent from agents
agents = agents[agents != agent_ego[0]]
agents_image = draw_boxes(self.raster_size, raster_from_world, np.append(agents, av_agent), 255)
ego_image = draw_boxes(self.raster_size, raster_from_world, agent_ego, 255)
# agents_images[..., i] = agents_image
# ego_images[..., i] = ego_image
# print(i)
if i == 0:
heads = np.expand_dims(agents_image + ego_image, axis=2)
agents_images[..., 0] += (agents_image * weights[i]).astype(np.float32)
ego_images[..., 0] += (ego_image * weights[i]).astype(np.float32)
# ego_heads =
# agents_images = np.expand_dims((weights * agents_images).max(2), axis=2)
# ego_images = np.expand_dims((weights * ego_images).max(2), axis=2)
# combine such that the image consists of [agent_t, agent_t-1, agent_t-2, ego_t, ego_t-1, ego_t-2]
out_im = np.concatenate((agents_images, ego_images, heads), -1)
# return out_im.astype(np.float32) / 255
return out_im.astype(np.float32) / np.amax(out_im)
def rasterize(self,
history_frames: np.ndarray,
history_agents: List[np.ndarray],
history_tl_faces: List[np.ndarray],
agent: Optional[np.ndarray] = None,
agents: Optional[np.ndarray] = None) -> np.ndarray:
frame = history_frames[0]
if agent is None:
translation = frame["ego_translation"]
yaw = frame["ego_rotation"]
else:
translation = agent["centroid"]
yaw = agent["yaw"]
self.agents = []
self.camera.position = np.array([translation[0], translation[1], 1.],
dtype=np.float32)
self.camera.rotation = np.array([0., 0., yaw], dtype=np.float32)
# Actor/Ego rasterizer
for i, (frame, agents) in enumerate(zip(history_frames,
history_agents)):
agents = filter_agents_by_labels(agents,
self.filter_agents_threshold)
# note the cast is for legacy support of dataset before April 2020
av_agent = get_ego_as_agent(frame).astype(agents.dtype)
if agent is None:
self.add_agents((agents, False))
self.add_agents((av_agent, True))
else:
agent_ego = filter_agents_by_track_id(agents,
agent["track_id"])
if len(agent_ego) == 0: # agent not in this history frame
self.add_agents((agents, False))
self.add_agents((av_agent, False))
else:
agents = agents[agents != agent_ego[0]]
self.add_agents((agents, False))
self.add_agents((av_agent, False))
self.add_agents((agent_ego, True))
# TODO: history frames.
break
glBindFramebuffer(GL_FRAMEBUFFER, self.output_fbo)
self.renderer.render([
(
'map_renderer',
[
# Map layer #1 - road surface
self.lane_surface_model,
# Map layer #2 - crosswalks,
self.crosswalk_model,
# Map layer #3 - lane lines
self.lane_lines_model,
]),
('entity_renderer', self.agents),
])
glPixelStorei(GL_PACK_ALIGNMENT, 1)
glReadBuffer(GL_COLOR_ATTACHMENT0)
image = glReadPixels(0, 0, self.raster_size[0], self.raster_size[1],
GL_RGB, GL_FLOAT)
glBindFramebuffer(GL_FRAMEBUFFER, 0)
return image
def rasterize(
self,
history_frames: np.ndarray,
history_agents: List[np.ndarray],
history_tl_faces: List[np.ndarray],
agent: Optional[np.ndarray] = None,
) -> np.ndarray:
# all frames are drawn relative to this one"
frame = history_frames[0]
if agent is None:
ego_translation_m = history_frames[0]["ego_translation"]
ego_yaw_rad = rotation33_as_yaw(frame["ego_rotation"])
else:
ego_translation_m = np.append(
agent["centroid"], history_frames[0]["ego_translation"][-1])
ego_yaw_rad = agent["yaw"]
raster_from_world = self.render_context.raster_from_world(
ego_translation_m, ego_yaw_rad)
# this ensures we always end up with fixed size arrays, +1 is because current time is also in the history
out_shape = (self.raster_size[1], self.raster_size[0],
self.history_num_frames + 1)
agents_images = np.zeros(out_shape, dtype=np.uint8)
ego_images = np.zeros(out_shape, dtype=np.uint8)
# --- 1. prepare agent keep indices for random agent drop augmentation ---
track_ids = np.concatenate([a["track_id"] for a in history_agents])
unique_track_ids = np.unique(track_ids).astype(np.int64)
n_max_agents = int(np.max(unique_track_ids) + 1) # +1 for host car.
unique_track_ids = np.concatenate([[0], unique_track_ids
]) # Add Host car, with id=0.
n_unique_agents = len(unique_track_ids)
# if not np.all(unique_track_ids == np.arange(np.max(unique_track_ids) + 1)):
# # It occured!! --> unique_track_ids is not continuous. Some numbers are filtered out.
# print("unique_track_ids", unique_track_ids, "is not continuous!!!")
if not self.eval and np.random.uniform() < self.agent_drop_ratio:
if self.agent_drop_prob < 0:
# Randomly decide number of agents to drop.
# 0 represents host car.
n_keep_agents = np.random.randint(0, n_unique_agents)
agent_keep_indices = np.random.choice(unique_track_ids,
n_keep_agents,
replace=False)
else:
# Decide agents to drop or not by agent_drop_prob.
agent_keep_indices = unique_track_ids[np.random.uniform(
0.0, 1.0, (n_unique_agents, )) > self.agent_drop_prob]
n_keep_agents = len(agent_keep_indices)
# Must keep ego agent!
if agent["track_id"] not in agent_keep_indices:
agent_keep_indices = np.append(agent_keep_indices,
agent["track_id"])
else:
n_keep_agents = n_unique_agents
# keep all agents
agent_keep_indices = None
# --- 2. prepare extent scale augmentation ratio ----
# TODO: create enough number of extent_ratio array. Actually n_keep_agents suffice but create n_max_agents
# for simplicity..
if self.eval:
# No augmentation.
agents_extent_ratio = np.ones((n_max_agents, 3))
elif self.min_extent_ratio == self.max_extent_ratio:
agents_extent_ratio = np.ones(
(n_max_agents, 3)) * self.min_extent_ratio
else:
agents_extent_ratio = np.random.uniform(self.min_extent_ratio,
self.max_extent_ratio,
(n_max_agents, 3))
ego_extent_ratio = agents_extent_ratio[0]
for i, (frame,
agents_) in enumerate(zip(history_frames, history_agents)):
agents = filter_agents_by_labels(agents_,
self.filter_agents_threshold)
if agent_keep_indices is not None:
# --- 1. apply agent drop augmentation ---
agents = agents[np.isin(agents["track_id"],
agent_keep_indices)]
# note the cast is for legacy support of dataset before April 2020
av_agent = get_ego_as_agent(frame).astype(agents.dtype)
# 2. --- apply extent scale augmentation ---
# TODO: Need to convert agents["track_id"] --> index based on `agent_keep_indices`,
# if we only create `agents_extent_ratio` of size `n_keep_agents`.
agents["extent"] *= agents_extent_ratio[agents["track_id"]]
av_agent[0]["extent"] *= ego_extent_ratio
if agent is None:
agents_image = draw_boxes(self.raster_size, raster_from_world,
agents, 255)
ego_image = draw_boxes(self.raster_size, raster_from_world,
av_agent, 255)
else:
agent_ego = filter_agents_by_track_id(agents,
agent["track_id"])
if agent_keep_indices is None or 0 in agent_keep_indices:
agents = np.append(agents, av_agent)
if len(agent_ego) == 0: # agent not in this history frame
agents_image = draw_boxes(self.raster_size,
raster_from_world, agents, 255)
ego_image = np.zeros_like(agents_image)
else: # add av to agents and remove the agent from agents
agents = agents[agents != agent_ego[0]]
agents_image = draw_boxes(self.raster_size,
raster_from_world, agents, 255)
ego_image = draw_boxes(self.raster_size, raster_from_world,
agent_ego, 255)
agents_images[..., i] = agents_image
ego_images[..., i] = ego_image
# combine such that the image consists of [agent_t, agent_t-1, agent_t-2, ego_t, ego_t-1, ego_t-2]
out_im = np.concatenate((agents_images, ego_images), -1)
return out_im.astype(np.float32) / 255
def _get_frame_data(mapAPI: MapAPI, frame: np.ndarray,
agents_frame: np.ndarray,
tls_frame: np.ndarray) -> FrameVisualization:
"""Get visualisation objects for the current frame.
:param mapAPI: mapAPI object (used for lanes, crosswalks etc..)
:param frame: the current frame (used for ego)
:param agents_frame: agents in this frame
:param tls_frame: the tls of this frame
:return: A FrameVisualization object. NOTE: trajectory are not included here
"""
ego_xy = frame["ego_translation"][:2]
#################
# plot lanes
lane_indices = indices_in_bounds(ego_xy,
mapAPI.bounds_info["lanes"]["bounds"], 50)
active_tl_ids = set(
filter_tl_faces_by_status(tls_frame, "ACTIVE")["face_id"].tolist())
lanes_vis: List[LaneVisualization] = []
for idx, lane_idx in enumerate(lane_indices):
lane_idx = mapAPI.bounds_info["lanes"]["ids"][lane_idx]
lane_tl_ids = set(mapAPI.get_lane_traffic_control_ids(lane_idx))
lane_colour = "gray"
for tl_id in lane_tl_ids.intersection(active_tl_ids):
lane_colour = COLORS[mapAPI.get_color_for_face(tl_id)]
lane_coords = mapAPI.get_lane_coords(lane_idx)
left_lane = lane_coords["xyz_left"][:, :2]
right_lane = lane_coords["xyz_right"][::-1, :2]
lanes_vis.append(
LaneVisualization(xs=np.hstack((left_lane[:, 0], right_lane[:,
0])),
ys=np.hstack((left_lane[:, 1], right_lane[:,
1])),
color=lane_colour))
#################
# plot crosswalks
crosswalk_indices = indices_in_bounds(
ego_xy, mapAPI.bounds_info["crosswalks"]["bounds"], 50)
crosswalks_vis: List[CWVisualization] = []
for idx in crosswalk_indices:
crosswalk = mapAPI.get_crosswalk_coords(
mapAPI.bounds_info["crosswalks"]["ids"][idx])
crosswalks_vis.append(
CWVisualization(xs=crosswalk["xyz"][:, 0],
ys=crosswalk["xyz"][:, 1],
color="yellow"))
#################
# plot ego and agents
agents_frame = np.insert(agents_frame, 0, get_ego_as_agent(frame))
box_world_coords = get_box_world_coords(agents_frame)
# ego
ego_vis = EgoVisualization(xs=box_world_coords[0, :, 0],
ys=box_world_coords[0, :, 1],
color="red",
center_x=agents_frame["centroid"][0, 0],
center_y=agents_frame["centroid"][0, 1])
# agents
agents_frame = agents_frame[1:]
box_world_coords = box_world_coords[1:]
agents_vis: List[AgentVisualization] = []
for agent, box_coord in zip(agents_frame, box_world_coords):
label_index = np.argmax(agent["label_probabilities"])
agent_type = PERCEPTION_LABELS[label_index]
agents_vis.append(
AgentVisualization(xs=box_coord[..., 0],
ys=box_coord[..., 1],
color="#1F77B4" if agent_type not in COLORS else
COLORS[agent_type],
track_id=agent["track_id"],
agent_type=PERCEPTION_LABELS[label_index],
prob=agent["label_probabilities"][label_index]))
return FrameVisualization(ego=ego_vis,
agents=agents_vis,
lanes=lanes_vis,
crosswalks=crosswalks_vis,
trajectories=[])
