def get_past_and_future_lanes(self, pos, quat, lane_data):
# TODO: filter conditions needs some work
if len(lane_data['incoming_lane_data']) == 0:
lane_data['incoming_lane_data'].append({'record': None})
if len(lane_data['outgoing_lane_data']) == 0:
lane_data['outgoing_lane_data'].append({'record': None})
if lane_data['closest_lane_data']['record'] is None or lane_data[
'incoming_lane_data'][0]['record'] is None or lane_data[
'outgoing_lane_data'][0]['record'] is None:
past_lane = np.array([pos[:2]])
future_lanes = [np.array([pos[:2]])]
else:
c_lane = lane_data['closest_lane_data']['poses']
p_lane = lane_data['incoming_lane_data'][0]['poses']
f_lanes = [l['poses'] for l in lane_data['outgoing_lane_data']]
c_lane_local = convert_global_coords_to_local(
c_lane[:, :2], pos, quat)
p_lane_local = convert_global_coords_to_local(
p_lane[:, :2], pos, quat)
f_lane_local = [
convert_global_coords_to_local(l[:, :2], pos, quat)
for l in f_lanes
]
cp_lane = np.vstack([p_lane_local, c_lane_local])
cpf_lanes = [np.vstack([cp_lane, f]) for f in f_lane_local]
past_lane_idx = np.argmax(cp_lane[:, 1] > 0)
if past_lane_idx == 0:
past_lane = np.array([pos[:2]])
future_lanes = [np.array([pos[:2]])]
else:
cp_lane_gb = convert_local_coords_to_global(cp_lane, pos, quat)
cpf_lanes_gb = [
convert_local_coords_to_global(l, pos, quat)
for l in cpf_lanes
]
past_lane = cp_lane_gb[:past_lane_idx, :]
future_lanes = [l[past_lane_idx:, :] for l in cpf_lanes_gb]
past_lane = process_to_len(past_lane,
500,
name='past_lane',
dim=0,
before_or_after='before',
mode='edge')
future_lanes = [
process_to_len(l,
500,
name='future_lane',
dim=0,
before_or_after='after',
mode='edge') for l in future_lanes
]
return past_lane, future_lanes
def plot_trajectory_distributions(self, ax, traj_dist_dict):
'''
traj_dist_dict = {
<instance_sample_token>: {"traj_dist": [[mean, cov], [], ...],
"frame": <"local" or "global">,
"pos": <np.ndarray or list>, # global coordinate of origin
"quat": <np.ndarray or list> # global rotation of origin
}
}
'''
for k, v in traj_dist_dict.items():
traj_dist = v['traj_dist']
for dist in traj_dist:
mean, cov = dist
if v['frame'] == 'local':
mean = convert_local_coords_to_global(
np.array(mean), v['pos'], v['quat'])
x, y = np.random.multivariate_normal(mean, cov, size=100).T
sns.kdeplot(x=x, y=y, cmap='Blues', ax=ax)
def test_heading_0(self):
rotation = (1, 0, 0, 0)
origin = (0, 0, 0)
offset = (50, 25, 0)
# Testing path along pos x direction
answer = convert_global_coords_to_local(self.along_pos_x, origin,
rotation)
np.testing.assert_allclose(answer, self.along_pos_y, atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, origin, rotation),
self.along_pos_x,
atol=1e-4)
answer = convert_global_coords_to_local(self.along_pos_x + [[50, 25]],
offset, rotation)
np.testing.assert_allclose(answer, self.along_pos_y, atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, offset, rotation),
self.along_pos_x + [[50, 25]],
atol=1e-4)
# Testing path along pos y direction
answer = convert_global_coords_to_local(self.along_pos_y, origin,
rotation)
np.testing.assert_allclose(answer, self.along_neg_x, atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, origin, rotation),
self.along_pos_y,
atol=1e-4)
answer = convert_global_coords_to_local(self.along_pos_y + [[50, 25]],
offset, rotation)
np.testing.assert_allclose(answer, self.along_neg_x, atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, offset, rotation),
self.along_pos_y + [[50, 25]],
atol=1e-4)
# Testing path along neg x direction
answer = convert_global_coords_to_local(self.along_neg_x, origin,
rotation)
np.testing.assert_allclose(answer, self.along_neg_y, atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, origin, rotation),
self.along_neg_x,
atol=1e-4)
answer = convert_global_coords_to_local(self.along_neg_x + [[50, 25]],
offset, rotation)
np.testing.assert_allclose(answer, self.along_neg_y, atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, offset, rotation),
self.along_neg_x + [[50, 25]],
atol=1e-4)
# Testing path along neg y direction
answer = convert_global_coords_to_local(self.along_neg_y, origin,
rotation)
np.testing.assert_allclose(answer, self.along_pos_x, atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, origin, rotation),
self.along_neg_y,
atol=1e-4)
answer = convert_global_coords_to_local(self.along_neg_y + [[50, 25]],
offset, rotation)
np.testing.assert_allclose(answer, self.along_pos_x, atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, offset, rotation),
self.along_neg_y + [[50, 25]],
atol=1e-4)
def test_heading_neg_pi_over_4(self):
rotation = (np.cos(-np.pi / 8), 0, 0, np.sin(-np.pi / 8))
origin = (0, 0, 0)
offset = (50, 25, 0)
# Testing path along pos x direction
answer = convert_global_coords_to_local(self.along_pos_x, origin,
rotation)
np.testing.assert_allclose(
answer,
self.y_equals_x *
[[-np.sqrt(2) / 2, np.sqrt(2) / 2]],
atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, origin, rotation),
self.along_pos_x,
atol=1e-4)
answer = convert_global_coords_to_local(self.along_pos_x + [[50, 25]],
offset, rotation)
np.testing.assert_allclose(
answer,
self.y_equals_x *
[[-np.sqrt(2) / 2, np.sqrt(2) / 2]],
atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, offset, rotation),
self.along_pos_x + [[50, 25]],
atol=1e-4)
# Testing path along pos y direction
answer = convert_global_coords_to_local(self.along_pos_y, origin,
rotation)
np.testing.assert_allclose(answer,
self.y_equals_x *
[[-np.sqrt(2) / 2, -np.sqrt(2) / 2]],
atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, origin, rotation),
self.along_pos_y,
atol=1e-4)
answer = convert_global_coords_to_local(self.along_pos_y + [[50, 25]],
offset, rotation)
np.testing.assert_allclose(answer,
self.y_equals_x *
[[-np.sqrt(2) / 2, -np.sqrt(2) / 2]],
atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, offset, rotation),
self.along_pos_y + [[50, 25]],
atol=1e-4)
# Testing path along neg x direction
answer = convert_global_coords_to_local(self.along_neg_x, origin,
rotation)
np.testing.assert_allclose(answer,
self.y_equals_x *
[[np.sqrt(2) / 2, -np.sqrt(2) / 2]],
atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, origin, rotation),
self.along_neg_x,
atol=1e-4)
answer = convert_global_coords_to_local(self.along_neg_x + [[50, 25]],
offset, rotation)
np.testing.assert_allclose(answer,
self.y_equals_x *
[[np.sqrt(2) / 2, -np.sqrt(2) / 2]],
atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, offset, rotation),
self.along_neg_x + [[50, 25]],
atol=1e-4)
# Testing path along neg y direction
answer = convert_global_coords_to_local(self.along_neg_y, origin,
rotation)
np.testing.assert_allclose(
answer,
self.y_equals_x * [[np.sqrt(2) / 2, np.sqrt(2) / 2]],
atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, origin, rotation),
self.along_neg_y,
atol=1e-4)
answer = convert_global_coords_to_local(self.along_neg_y + [[50, 25]],
offset, rotation)
np.testing.assert_allclose(
answer,
self.y_equals_x * [[np.sqrt(2) / 2, np.sqrt(2) / 2]],
atol=1e-4)
np.testing.assert_allclose(convert_local_coords_to_global(
answer, offset, rotation),
self.along_neg_y + [[50, 25]],
atol=1e-4)
def plot_car_to_predict(self,
ax,
agent_future: np.ndarray,
agent_past: np.ndarray,
instance_token: str,
sample_token: str,
text_box: bool = True,
predictions: list = None,
agent_state_dict: dict = None):
'''
predictions = {
'name': {'data': <data>, 'color': <coloar>, 'frame': <frame>, 'style':'.'}
}
'''
## plot car ####
ann = self.helper.get_sample_annotation(instance_token, sample_token)
category = ann['category_name']
if len(ann['attribute_tokens']) != 0:
attribute = self.nusc.get('attribute',
ann['attribute_tokens'][0])['name']
else:
attribute = ""
agent_yaw = Quaternion(ann['rotation'])
agent_yaw = quaternion_yaw(agent_yaw)
agent_yaw = angle_of_rotation(agent_yaw)
agent_yaw = np.rad2deg(agent_yaw)
self.plot_elements([ann['translation'][0], ann['translation'][1]],
agent_yaw, 'current_car', ax)
if text_box:
self.plot_text_box(
ax, category,
[ann['translation'][0] + 1.2, ann['translation'][1]])
self.plot_text_box(
ax, attribute,
[ann['translation'][0] + 1.2, ann['translation'][1] + 1.2])
if agent_state_dict is not None:
state_str = ""
for k, v in agent_state_dict.items():
state_str += f"{k[0]}:{v:.2f}, "
self.plot_text_box(
ax, state_str,
[ann['translation'][0] + 1.2, ann['translation'][1] + 3.2])
agent_pos = [ann['translation'][0], ann['translation'][1]]
agent_yaw_deg = agent_yaw
# plot ground truth
if len(agent_future) > 0:
ax.scatter(agent_future[:, 0],
agent_future[:, 1],
s=20,
c='yellow',
alpha=1.0,
zorder=200)
if len(agent_past) > 0:
ax.scatter(agent_past[:, 0],
agent_past[:, 1],
s=20,
c='k',
alpha=0.5,
zorder=200)
# plot predictions
if predictions is not None:
for k, v in viewitems(predictions):
if v['frame'] == 'local':
v['data'] = convert_local_coords_to_global(
v['data'], np.array(agent_pos),
np.array(ann['rotation']))
if 'style' not in v.keys():
v['style'] = '.'
if v['style'] == '.':
ax.scatter(v['data'][:, 0],
v['data'][:, 1],
s=20,
c=v['color'],
alpha=1.0,
zorder=2)
elif v['style'] == '-':
ax.plot(v['data'][:, 0],
v['data'][:, 1],
c=v['color'],
alpha=1.0,
zorder=2)
else:
raise ValueError('style not supported')
return agent_pos, ann['rotation']