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 test_draw_boxes() -> None:
agents = np.zeros(2, dtype=AGENT_DTYPE)
agents[0]["extent"] = (20, 20, 20)
agents[0]["centroid"] = (100, 100)
agents[1]["extent"] = (20, 20, 20)
agents[1]["centroid"] = (150, 150)
to_image_space = np.eye(3)
im = draw_boxes((200, 200), to_image_space, agents, color=1)
assert im.sum() == (21 * 21) * 2
def test_draw_boxes() -> None:
centroid_1 = (90, 100)
centroid_2 = (150, 160)
agents = np.zeros(2, dtype=AGENT_DTYPE)
agents[0]["extent"] = (20, 20, 20)
agents[0]["centroid"] = centroid_1
agents[1]["extent"] = (20, 20, 20)
agents[1]["centroid"] = centroid_2
to_image_space = np.eye(3)
im = draw_boxes((200, 200), to_image_space, agents, color=1)
# due to subpixel precision we can't check the exact number of pixels
# check that a 10x10 centred on the boxes is all 1
assert np.allclose(im[centroid_1[1] - 5: centroid_1[1] + 5, centroid_1[0] - 5: centroid_1[0] + 5], 1)
assert np.allclose(im[centroid_2[1] - 5: centroid_2[1] + 5, centroid_2[0] - 5: centroid_2[0] + 5], 1)
def test_empty_boxes() -> None:
# naive test with empty arrays
agents = np.empty(0, dtype=AGENT_DTYPE)
to_image_space = np.eye(3)
im = draw_boxes((200, 200), to_image_space, agents, color=255)
assert im.sum() == 0
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,
) -> 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
