def get_track_box(annotation: Dict[str, Any],
center_coordinates: Tuple[float, float],
center_pixels: Tuple[float, float],
resolution: float = 0.1) -> np.ndarray:
"""
Get four corners of bounding box for agent in pixels.
:param annotation: The annotation record of the agent.
:param center_coordinates: (x, y) coordinates in global frame
of the center of the image.
:param center_pixels: (row_index, column_index) location of the center
of the image in pixel coordinates.
:param resolution: Resolution pixels/meter of the image.
"""
assert resolution > 0
location = annotation['translation'][:2]
yaw_in_radians = quaternion_yaw(Quaternion(annotation['rotation']))
row_pixel, column_pixel = convert_to_pixel_coords(location,
center_coordinates,
center_pixels,
resolution)
width = annotation['size'][0] / resolution
length = annotation['size'][1] / resolution
# Width and length are switched here so that we can draw them along the x-axis as
# opposed to the y. This makes rotation easier.
return pixels_to_box_corners(row_pixel, column_pixel, length, width,
yaw_in_radians)
def make_representation(self, instance_token: str,
sample_token: str) -> np.ndarray:
"""
Draws agent boxes with faded history into a black background.
:param instance_token: Instance token.
:param sample_token: Sample token.
:return: np.ndarray representing a 3 channel image.
"""
# Taking radius around track before to ensure all actors are in image
buffer = max([
self.meters_ahead, self.meters_behind, self.meters_left,
self.meters_right
]) * 2
image_side_length = int(buffer / self.resolution)
# We will center the track in the image
central_track_pixels = (image_side_length / 2, image_side_length / 2)
base_image = np.zeros((image_side_length, image_side_length, 3))
history = self.helper.get_past_for_sample(sample_token,
self.seconds_of_history,
in_agent_frame=False,
just_xy=False)
history = reverse_history(history)
present_time = self.helper.get_annotations_for_sample(sample_token)
history = add_present_time_to_history(present_time, history)
center_agent_annotation = self.helper.get_sample_annotation(
instance_token, sample_token)
draw_agent_boxes(center_agent_annotation,
central_track_pixels,
history,
base_image,
resolution=self.resolution,
get_color=self.color_mapping)
center_agent_yaw = quaternion_yaw(
Quaternion(center_agent_annotation['rotation']))
rotation_mat = get_rotation_matrix(base_image.shape, center_agent_yaw)
rotated_image = cv2.warpAffine(
base_image, rotation_mat,
(base_image.shape[1], base_image.shape[0]))
row_crop, col_crop = get_crops(self.meters_ahead, self.meters_behind,
self.meters_left, self.meters_right,
self.resolution, image_side_length)
return rotated_image[row_crop, col_crop].astype('uint8')
def make_representation(self, instance_token: str,
sample_token: str) -> np.ndarray:
"""
Represents agents as a 4 channel image. Channel 0: 1 if an agent is present at this location 0 otherwise
Channel 1, 2, 3 velocity, acc and yaw-rate of agent
If multiple agents are present at same location for given resolution, channel 0 is summed, min of other
channels retained. (Probably need a better solution than min). All values are scalars, direction can be
inferred using lane direction from static layer.
:param instance_token: Instance token.
:param sample_token: Sample token.
:return: np.ndarray representing a 4 channel image.
"""
# Taking radius around track before to ensure all actors are in image
buffer = max([
self.meters_ahead, self.meters_behind, self.meters_left,
self.meters_right
]) * 2
image_side_length = int(buffer / self.resolution)
# We will center the track in the image
central_track_pixels = (image_side_length / 2, image_side_length / 2)
base_image = np.zeros((image_side_length, image_side_length, 4))
annotations = self.helper.get_annotations_for_sample(sample_token)
center_agent_annotation = self.helper.get_sample_annotation(
instance_token, sample_token)
populate_agent_states(center_agent_annotation,
central_track_pixels,
annotations,
base_image,
self.helper,
resolution=self.resolution)
# Rotate and crop representation:
center_agent_yaw = quaternion_yaw(
Quaternion(center_agent_annotation['rotation']))
rotation_mat = get_rotation_matrix(base_image.shape, center_agent_yaw)
rotated_image = cv2.warpAffine(
base_image, rotation_mat,
(base_image.shape[1], base_image.shape[0]))
row_crop, col_crop = get_crops(self.meters_ahead, self.meters_behind,
self.meters_left, self.meters_right,
self.resolution, image_side_length)
return rotated_image[row_crop, col_crop]
def generate_mask(self, translation, rotation, sample_token: str, seconds=3, show_agent=True):
buffer = max([self.meters_ahead, self.meters_behind, self.meters_left, self.meters_right]) * 2
image_side_length = int(buffer / self.resolution)
central_track_pixels = (image_side_length / 2, image_side_length / 2)
base_image = np.zeros((image_side_length, image_side_length, 3))
present_time_annotation = self.helper.get_annotations_for_sample(sample_token)
agent_x, agent_y = translation[:2]
translation_xy = []
past_xy = []
future_xy = []
for annotation in present_time_annotation:
past_xy_temp = self.helper.get_past_for_agent(
annotation['instance_token'], sample_token, seconds=seconds, in_agent_frame=False)
future_xy_temp = self.helper.get_future_for_agent(
annotation['instance_token'], sample_token, seconds=seconds, in_agent_frame=False)
# vel = self.helper.get_velocity_for_agent(annotation['instance_token'], sample_token)
# accel = self.helper.get_acceleration_for_agent(annotation['instance_token'], sample_token)
# yaw_rate = self.helper.get_heading_change_rate_for_agent(annotation['instance_token'], sample_token)
if annotation['category_name'].split('.')[0] != 'vehicle':
continue
translation_xy.append(annotation['translation'][:2])
past_xy.append(past_xy_temp)
future_xy.append(future_xy_temp)
if show_agent:
box = get_track_box(annotation, (agent_x, agent_y), central_track_pixels, self.resolution)
color = self.color_mapping(annotation['category_name'])
cv2.fillPoly(base_image, pts=[np.int0(box)], color=color)
xy_global = [np.asarray(past_xy), np.asarray(future_xy), np.asarray(translation_xy)]
center_agent_yaw = quaternion_yaw(Quaternion(rotation))
rotation_mat = get_rotation_matrix(base_image.shape, center_agent_yaw)
rotated_image = cv2.warpAffine(base_image, rotation_mat, (base_image.shape[1],
base_image.shape[0]))
row_crop, col_crop = get_crops(self.meters_ahead, self.meters_behind,
self.meters_left, self.meters_right, self.resolution,
image_side_length)
return rotated_image[row_crop, col_crop].astype('uint8'), xy_global
def generate_mask(self, translation, rotation, sample_token: str):
buffer = max([
self.meters_ahead, self.meters_behind, self.meters_left,
self.meters_right
]) * 2
image_side_length = int(buffer / self.resolution)
central_track_pixels = (image_side_length / 2, image_side_length / 2)
base_image = np.zeros((image_side_length, image_side_length, 3))
history = self.helper.get_past_for_sample(sample_token,
self.seconds_of_history,
in_agent_frame=False,
just_xy=False)
history = reverse_history(history)
present_time = self.helper.get_annotations_for_sample(sample_token)
history = add_present_time_to_history(present_time, history)
agent_x, agent_y = translation[:2]
for ins_token, annot in history.items():
num_points = len(annot)
for i, annotation in enumerate(annot):
box = get_track_box(annotation, (agent_x, agent_y),
central_track_pixels, self.resolution)
color = self.color_mapping(annotation['category_name'])
# Don't fade the colors if there is no history
if num_points > 1:
color = fade_color(color, i, num_points - 1)
cv2.fillPoly(base_image, pts=[np.int0(box)], color=color)
center_agent_yaw = quaternion_yaw(Quaternion(rotation))
rotation_mat = get_rotation_matrix(base_image.shape, center_agent_yaw)
rotated_image = cv2.warpAffine(
base_image, rotation_mat,
(base_image.shape[1], base_image.shape[0]))
row_crop, col_crop = get_crops(self.meters_ahead, self.meters_behind,
self.meters_left, self.meters_right,
self.resolution, image_side_length)
return rotated_image[row_crop, col_crop].astype('uint8')
def make_representation(self,
instance_token: str = None,
sample_token: str = None,
ego: bool = False,
ego_pose=None,
include_history=True) -> np.ndarray:
"""
Draws agent boxes with faded history into a black background.
:param instance_token: Instance token.
:param sample_token: Sample token.
:return: np.ndarray representing a 3 channel image.
"""
# Taking radius around track before to ensure all actors are in image
buffer = max([
self.meters_ahead, self.meters_behind, self.meters_left,
self.meters_right
]) * 2
image_side_length = int(buffer / self.resolution)
# We will center the track in the image
central_track_pixels = (image_side_length / 2, image_side_length / 2)
base_image = np.zeros((image_side_length, image_side_length, 3))
history = self.helper.get_past_for_sample(sample_token,
self.seconds_of_history,
in_agent_frame=False,
just_xy=False)
history = reverse_history(history)
present_time = self.helper.get_annotations_for_sample(sample_token)
history = add_present_time_to_history(present_time, history)
if not ego:
center_agent_annotation = self.helper.get_sample_annotation(
instance_token, sample_token)
else:
sample_ = self.helper.data.get('sample', sample_token)
sample_data = self.helper.data.get('sample_data',
sample_['data']['CAM_FRONT'])
if ego_pose is None:
ego_pose = self.helper.data.get('ego_pose',
sample_data['ego_pose_token'])
center_agent_annotation = {
'translation': ego_pose['translation'],
'rotation': ego_pose['rotation'],
'instance_token': None
}
if not include_history:
curr = {}
for token, anns in history.items():
curr[token] = [anns[-1]]
history = curr
draw_agent_boxes(center_agent_annotation,
central_track_pixels,
history,
base_image,
resolution=self.resolution,
get_color=self.color_mapping,
ego=ego,
ego_pose=ego_pose)
center_agent_yaw = quaternion_yaw(
Quaternion(center_agent_annotation['rotation']))
rotation_mat = get_rotation_matrix(base_image.shape, center_agent_yaw)
rotated_image = cv2.warpAffine(
base_image, rotation_mat,
(base_image.shape[1], base_image.shape[0]))
row_crop, col_crop = get_crops(self.meters_ahead, self.meters_behind,
self.meters_left, self.meters_right,
self.resolution, image_side_length)
return rotated_image[row_crop, col_crop].astype('uint8')
def generate_virtual_mask(self, translation, rotation, lanes, sample_token: str, show_agent=True,
past_trj_len=4, future_trj_len=6, min_dist=6):
buffer = max([self.meters_ahead, self.meters_behind, self.meters_left, self.meters_right]) * 2
image_side_length = int(buffer / self.resolution)
central_track_pixels = (image_side_length / 2, image_side_length / 2)
base_image = np.zeros((image_side_length, image_side_length, 3))
translation_xy = []
past_xy = []
future_xy = []
for lane in lanes:
length = len(lane)
if length < (past_trj_len + future_trj_len + 1):
continue
location = None
current_index = None
for _ in range(5):
is_collide = False
if (length - past_trj_len - future_trj_len - 1) < 1:
continue
# start_index = np.random.randint(low=0, high=(length - past_trj_len - future_trj_len - 1))
start_index = 0
max_gap = (length - start_index) // (past_trj_len + future_trj_len + 1)
if max_gap < 2:
continue
# gap = np.random.randint(low=1, high=max_gap+1)
gap = max_gap
# current_index = start_index + (gap * past_trj_len)
# location_ = lane[current_index, :2]
mask = np.arange(start_index, length, step=gap)
lane_ = lane[mask]
current_index = past_trj_len
location_ = lane_[current_index, :2]
# check collision
for xy in translation_xy:
diff = location_ - xy
if np.linalg.norm(diff) < min_dist:
is_collide = True
break
if not is_collide:
location = location_
lane = lane_
break
if location is None:
continue
translation_xy.append(location)
past_xy.append(lane[current_index - past_trj_len:current_index])
future_xy.append(lane[current_index + 1:current_index + future_trj_len + 1])
# draw virtual agent mask
if show_agent:
agent_x, agent_y = translation[:2]
ax, ay = lane[current_index, 0], lane[current_index, 1]
bx, by = lane[current_index + 1, 0], lane[current_index + 1, 1]
if ax == bx:
bx += 0.0001
rot = np.arctan((by - ay) / (bx - ax))
yaw_in_radians = quaternion_yaw(Quaternion(axis=[0, 0, 1], angle=rot))
row_pixel, column_pixel = convert_to_pixel_coords(
location, (agent_x, agent_y), central_track_pixels, self.resolution)
size = [2.2, 5.6, 2.1]
width = size[0] / self.resolution
length = size[1] / self.resolution
box = pixels_to_box_corners(row_pixel, column_pixel, length, width, yaw_in_radians)
color = (255, 255, 0)
cv2.fillPoly(base_image, pts=[np.int0(box)], color=color)
xy_global = [np.asarray(past_xy), np.asarray(future_xy), np.asarray(translation_xy)]
center_agent_yaw = quaternion_yaw(Quaternion(rotation))
rotation_mat = get_rotation_matrix(base_image.shape, center_agent_yaw)
rotated_image = cv2.warpAffine(base_image, rotation_mat, (base_image.shape[1], base_image.shape[0]))
row_crop, col_crop = get_crops(self.meters_ahead, self.meters_behind,
self.meters_left, self.meters_right, self.resolution,
image_side_length)
return rotated_image[row_crop, col_crop].astype('uint8'), xy_global
def generate_mask_with_path(self,
instance_token: str,
sample_token: str,
seconds=3,
vehicle_only=True):
annotation = self.helper.get_sample_annotation(
instance_token, sample_token) # agent annotation
ego_pose_xy = annotation['translation']
ego_pose_rotation = annotation['rotation']
buffer = max([
self.meters_ahead, self.meters_behind, self.meters_left,
self.meters_right
]) * 2
image_side_length = int(buffer / self.resolution)
central_track_pixels = (image_side_length / 2, image_side_length / 2)
base_image = np.zeros((image_side_length, image_side_length, 3))
present_time_annotation = self.helper.get_annotations_for_sample(
sample_token)
instance_tokens = []
category_names = []
translations = []
pasts = []
futures = []
velocities = []
accelerations = []
yaw_rates = []
instance_idx = -1
for idx, annotation in enumerate(present_time_annotation):
if vehicle_only and annotation['category_name'].split(
'.')[0] != 'vehicle':
continue
elif annotation['instance_token'] == instance_token:
instance_idx = len(instance_tokens)
past_xy_temp = self.helper.get_past_for_agent(
annotation['instance_token'],
sample_token,
seconds=seconds,
in_agent_frame=False)
future_xy_temp = self.helper.get_future_for_agent(
annotation['instance_token'],
sample_token,
seconds=seconds,
in_agent_frame=False)
instance_tokens.append(annotation['instance_token'])
category_names.append(annotation['category_name'])
translations.append(annotation['translation'][:2])
pasts.append(np.flip(past_xy_temp, axis=0).tolist())
futures.append(future_xy_temp.tolist())
velocities.append(
self.helper.get_velocity_for_agent(
annotation['instance_token'], sample_token))
accelerations.append(
self.helper.get_acceleration_for_agent(
annotation['instance_token'], sample_token))
yaw_rates.append(
self.helper.get_heading_change_rate_for_agent(
annotation['instance_token'], sample_token))
box = get_track_box(annotation, (ego_pose_xy[0], ego_pose_xy[1]),
central_track_pixels, self.resolution)
color = self.color_mapping(annotation['category_name'])
cv2.fillPoly(base_image, pts=[np.int0(box)], color=color)
total_agents_annotation = {
'instance_tokens': instance_tokens,
'category_names': category_names,
'translations': translations,
'pasts': pasts,
'futures': futures,
'velocities': velocities,
'accelerations': accelerations,
'yaw_rates': yaw_rates
}
agent_annotation = {
'category_name': category_names[instance_idx],
'translation': translations[instance_idx],
'past': pasts[instance_idx],
'future': futures[instance_idx],
'velocity': velocities[instance_idx],
'acceleration': accelerations[instance_idx],
'yaw_rate': yaw_rates[instance_idx]
}
center_agent_yaw = quaternion_yaw(Quaternion(ego_pose_rotation))
rotation_mat = get_rotation_matrix(base_image.shape, center_agent_yaw)
rotated_image = cv2.warpAffine(
base_image, rotation_mat,
(base_image.shape[1], base_image.shape[0]))
row_crop, col_crop = get_crops(self.meters_ahead, self.meters_behind,
self.meters_left, self.meters_right,
self.resolution, image_side_length)
agent_mask = rotated_image[row_crop, col_crop].astype('uint8')
return agent_mask, total_agents_annotation, agent_annotation