def verify_point_in_polygon_for_lanes():
"""
"""
avm = ArgoverseMap()
# ref_query_x = 422.
# ref_query_y = 1005.
ref_query_x = -662
ref_query_y = 2817
city_name = "MIA"
for trial_idx in range(10):
query_x = ref_query_x + (np.random.rand() - 0.5) * 10
query_y = ref_query_y + (np.random.rand() - 0.5) * 10
fig = plt.figure(figsize=(22.5, 8))
ax = fig.add_subplot(111)
ax.scatter([query_x], [query_y], 100, color="k", marker=".")
occupied_lane_ids = avm.get_lane_segments_containing_xy(
query_x, query_y, city_name)
for occupied_lane_id in occupied_lane_ids:
halluc_lane_polygon = avm.get_lane_segment_polygon(
occupied_lane_id, city_name)
plot_lane_segment_patch(halluc_lane_polygon,
ax,
color="y",
alpha=0.3)
nearby_lane_ids = avm.get_lane_ids_in_xy_bbox(query_x, query_y,
city_name)
nearby_lane_ids = set(nearby_lane_ids) - set(occupied_lane_ids)
for nearby_lane_id in nearby_lane_ids:
halluc_lane_polygon = avm.get_lane_segment_polygon(
nearby_lane_id, city_name)
plot_lane_segment_patch(halluc_lane_polygon,
ax,
color="r",
alpha=0.3)
ax.axis("equal")
plt.show()
plt.close("all")
def get_point_in_polygon_score(self, lane_seq: List[int],
xy_seq: np.ndarray, city_name: str,
avm: ArgoverseMap) -> int:
"""Get the number of coordinates that lie insde the lane seq polygon.
Args:
lane_seq: Sequence of lane ids
xy_seq: Trajectory coordinates
city_name: City name (PITT/MIA)
avm: Argoverse map_api instance
Returns:
point_in_polygon_score: Number of coordinates in the trajectory that lie within the lane sequence
"""
lane_seq_polygon = cascaded_union([
Polygon(avm.get_lane_segment_polygon(lane, city_name)).buffer(0)
for lane in lane_seq
])
point_in_polygon_score = 0
for xy in xy_seq:
point_in_polygon_score += lane_seq_polygon.contains(Point(xy))
return point_in_polygon_score
def plot_nearby_halluc_lanes(
ax: plt.Axes,
city_name: str,
avm: ArgoverseMap,
query_x: float,
query_y: float,
patch_color: str = "r",
radius: float = 20.0,
) -> None:
"""Produce lane segment graphs for visual verification."""
nearby_lane_ids = avm.get_lane_ids_in_xy_bbox(query_x, query_y, city_name,
radius)
for nearby_lane_id in nearby_lane_ids:
halluc_lane_polygon = avm.get_lane_segment_polygon(
nearby_lane_id, city_name)
plot_lane_segment_patch(halluc_lane_polygon,
ax,
color=patch_color,
alpha=0.3)
plt.text(
halluc_lane_polygon[:, 0].mean(),
halluc_lane_polygon[:, 1].mean(),
str(nearby_lane_id),
)
class argoverse_processor(object):
"""
This object can extract context information from raw map file and raw trajectory file.
"""
def __init__(self, scenario_path = _sample_path):
self.map_pres = ArgoverseMap()
self.scenarios = ArgoverseForecastingLoader(scenario_path)
self.seq_lane_props = self.map_pres.city_lane_centerlines_dict
def acquire_agent_trajectory(self, one_scenario):
# There is 5 seconds of one scenario in total. 0.1s interval so that it has 50 items.
return one_scenario.agent_traj
def acquire_city(self, one_scenario):
return one_scenario.seq_df['CITY_NAME'].values[0]
def acquire_lane_centerlines_and_edgelines(self, one_scenario):
city_name = self.acquire_city(one_scenario)
seq_def = one_scenario.seq_df
seq_lanes = self.seq_lane_props[city_name]
x_min = min(seq_def['X'])
x_max = max(seq_def['X'])
y_min = min(seq_def['Y'])
y_max = max(seq_def['Y'])
lane_centerlines = []
lane_ids = []
# Get lane centerlines which lie within the range of trajectories
for lane_id, lane_props in seq_lanes.items():
lane_cl = lane_props.centerline
if (
np.min(lane_cl[:, 0]) < x_max
and np.min(lane_cl[:, 1]) < y_max
and np.max(lane_cl[:, 0]) > x_min
and np.max(lane_cl[:, 1]) > y_min
):
lane_centerlines.append(lane_cl)
lane_ids.append(lane_id)
edges1 = []
edges2 = []
for lane_id in lane_ids:
lane_polyline = self.map_pres.get_lane_segment_polygon(lane_id, city_name)
lane_polyline = lane_polyline[:-1]
break_point = len(lane_polyline)/2
break_point = int(break_point)
edges1.append(lane_polyline[:break_point])
edges2.append(lane_polyline[break_point:])
return lane_centerlines, edges1, edges2
def visualization_lanes(self, one_scenario):
plt.figure(figsize=(8, 7))
lane_cls, lane_eds1, lane_eds2 = self.acquire_lane_centerlines_and_edgelines(one_scenario)
for lane_cl in lane_cls:
a, = plt.plot(lane_cl[:, 0], lane_cl[:, 1], '--', label='lane_centerline', color='#17A9C3', linewidth=1) # color='#C7398D')
for lane_ed in lane_eds1:
b, = plt.plot(lane_ed[:, 0], lane_ed[:, 1], label='lane_edgeline', color='#C7398D', linewidth=1) # color='#314E87')
for lane_ed in lane_eds2:
c, = plt.plot(lane_ed[:, 0], lane_ed[:, 1], label='lane_edgeline', color='#C7398D', linewidth=1)
plt.legend([a, b], ['lane_centerline', 'lane_edgeline'], fontsize=10)
plt.axis('off')
plt.show()
def visualization_trajectory(self, one_scenario):
plt.figure(figsize=(8, 7))
lane_cls, lane_eds1, lane_eds2 = self.acquire_lane_centerlines_and_edgelines(one_scenario)
for lane_cl in lane_cls:
a, = plt.plot(lane_cl[:, 0], lane_cl[:, 1], '--', color='grey', linewidth=1) # color='#C7398D')
for lane_ed in lane_eds1:
b, = plt.plot(lane_ed[:, 0], lane_ed[:, 1], color='grey', linewidth=1) # color='#314E87')
for lane_ed in lane_eds2:
c, = plt.plot(lane_ed[:, 0], lane_ed[:, 1], color='grey', linewidth=1)
traj = self.acquire_agent_trajectory(one_scenario)
train_traj = traj[:20]
pred_traj = traj[20:]
d, = plt.plot(train_traj[:, 0], train_traj[:, 1], '-', color='#17A9C3', alpha=1, linewidth=2, zorder=2)
e, = plt.plot(train_traj[-1, 0], train_traj[-1, 1], 'o', color='#17A9C3', alpha=1, markersize=7, zorder=2)
f, = plt.plot(pred_traj[:, 0], pred_traj[:, 1], '-', color='#C7398D', alpha=1, linewidth=1, zorder=2)
plt.legend([a, b, d, f], ['lane_centerline', 'lane_edgeline', 'observed_agent_trajectory', 'future_agent_trajectory'], fontsize=10)
plt.axis('off')
plt.show()
if len(candidate_cl)<=2:
continue
label = int(label)
s_line = source_closest_line.centerline
t_line = target_closest_line.centerline
label_line = candidate_cl[label]
other_lines = candidate_cl[0:label] + candidate_cl[label+1:]
intersections = []
for row in separation_points:
intersection = Polygon(avm.get_lane_segment_polygon(row["ids"][0], city)[:, 0:2])
for id in row["ids"][1:]:
intersection = intersection.intersection(Polygon(avm.get_lane_segment_polygon(id, city)[:, 0:2]))
intersections.append(intersection)
intersections = [intersection.exterior.coords for intersection in intersections]
points = [ row["coord"] for row in separation_points]
image = draw(source, target, source_ref_line, label_line=label_line, other_lines=other_lines, points=points, polygons=intersections)
cv2.imwrite(output_file%(i+1), image)
class ArgoDataset(Dataset):
def __init__(self, split, config, train=True):
self.config = config
self.train = train
if 'preprocess' in config and config['preprocess']:
if train:
self.split = np.load(self.config['preprocess_train'],
allow_pickle=True)
else:
self.split = np.load(self.config['preprocess_val'],
allow_pickle=True)
else:
self.avl = ArgoverseForecastingLoader(split)
self.avl.seq_list = sorted(self.avl.seq_list)
self.am = ArgoverseMap()
if 'raster' in config and config['raster']:
#TODO: DELETE
self.map_query = MapQuery(config['map_scale'])
def __getitem__(self, idx):
if 'preprocess' in self.config and self.config['preprocess']:
data = self.split[idx]
if self.train and self.config['rot_aug']:
new_data = dict()
for key in ['city', 'orig', 'gt_preds', 'has_preds']:
if key in data:
new_data[key] = ref_copy(data[key])
dt = np.random.rand() * self.config['rot_size'] #np.pi * 2.0
theta = data['theta'] + dt
new_data['theta'] = theta
new_data['rot'] = np.asarray(
[[np.cos(theta), -np.sin(theta)],
[np.sin(theta), np.cos(theta)]], np.float32)
rot = np.asarray([[np.cos(-dt), -np.sin(-dt)],
[np.sin(-dt), np.cos(-dt)]], np.float32)
new_data['feats'] = data['feats'].copy()
new_data['feats'][:, :, :2] = np.matmul(
new_data['feats'][:, :, :2], rot)
new_data['ctrs'] = np.matmul(data['ctrs'], rot)
graph = dict()
for key in [
'num_nodes', 'turn', 'control', 'intersect', 'pre',
'suc', 'lane_idcs', 'left_pairs', 'right_pairs',
'left', 'right'
]:
graph[key] = ref_copy(data['graph'][key])
graph['ctrs'] = np.matmul(data['graph']['ctrs'], rot)
graph['feats'] = np.matmul(data['graph']['feats'], rot)
new_data['graph'] = graph
data = new_data
else:
new_data = dict()
for key in [
'city', 'orig', 'gt_preds', 'has_preds', 'theta',
'rot', 'feats', 'ctrs', 'graph'
]:
if key in data:
new_data[key] = ref_copy(data[key])
data = new_data
if 'raster' in self.config and self.config['raster']:
data.pop('graph')
x_min, x_max, y_min, y_max = self.config['pred_range']
cx, cy = data['orig']
region = [cx + x_min, cx + x_max, cy + y_min, cy + y_max]
raster = self.map_query.query(region, data['theta'],
data['city'])
data['raster'] = raster
return data
data = self.read_argo_data(idx)
data = self.get_obj_feats(data)
data['idx'] = idx
if 'raster' in self.config and self.config['raster']:
x_min, x_max, y_min, y_max = self.config['pred_range']
cx, cy = data['orig']
region = [cx + x_min, cx + x_max, cy + y_min, cy + y_max]
raster = self.map_query.query(region, data['theta'], data['city'])
data['raster'] = raster
return data
data['graph'] = self.get_lane_graph(data)
return data
def __len__(self):
if 'preprocess' in self.config and self.config['preprocess']:
return len(self.split)
else:
return len(self.avl)
def read_argo_data(self, idx):
city = copy.deepcopy(self.avl[idx].city)
"""TIMESTAMP,TRACK_ID,OBJECT_TYPE,X,Y,CITY_NAME"""
df = copy.deepcopy(self.avl[idx].seq_df)
agt_ts = np.sort(np.unique(df['TIMESTAMP'].values))
mapping = dict()
for i, ts in enumerate(agt_ts):
mapping[ts] = i
trajs = np.concatenate(
(df.X.to_numpy().reshape(-1, 1), df.Y.to_numpy().reshape(-1, 1)),
1)
steps = [mapping[x] for x in df['TIMESTAMP'].values]
steps = np.asarray(steps, np.int64)
objs = df.groupby(['TRACK_ID', 'OBJECT_TYPE']).groups
keys = list(objs.keys())
obj_type = [x[1] for x in keys]
agt_idx = obj_type.index('AGENT')
idcs = objs[keys[agt_idx]]
agt_traj = trajs[idcs]
agt_step = steps[idcs]
del keys[agt_idx]
ctx_trajs, ctx_steps = [], []
for key in keys:
idcs = objs[key]
ctx_trajs.append(trajs[idcs])
ctx_steps.append(steps[idcs])
data = dict()
data['city'] = city
data['trajs'] = [agt_traj] + ctx_trajs
data['steps'] = [agt_step] + ctx_steps
return data
def get_obj_feats(self, data):
orig = data['trajs'][0][19].copy().astype(np.float32)
if self.train and self.config['rot_aug']:
theta = np.random.rand() * np.pi * 2.0
else:
pre = data['trajs'][0][18] - orig
theta = np.pi - np.arctan2(pre[1], pre[0])
rot = np.asarray([[np.cos(theta), -np.sin(theta)],
[np.sin(theta), np.cos(theta)]], np.float32)
feats, ctrs, gt_preds, has_preds = [], [], [], []
for traj, step in zip(data['trajs'], data['steps']):
if 19 not in step:
continue
gt_pred = np.zeros((30, 2), np.float32)
has_pred = np.zeros(30, np.bool)
future_mask = np.logical_and(step >= 20, step < 50)
post_step = step[future_mask] - 20
post_traj = traj[future_mask]
gt_pred[post_step] = post_traj
has_pred[post_step] = 1
obs_mask = step < 20
step = step[obs_mask]
traj = traj[obs_mask]
idcs = step.argsort()
step = step[idcs]
traj = traj[idcs]
for i in range(len(step)):
if step[i] == 19 - (len(step) - 1) + i:
break
step = step[i:]
traj = traj[i:]
feat = np.zeros((20, 3), np.float32)
feat[step, :2] = np.matmul(rot, (traj - orig.reshape(-1, 2)).T).T
feat[step, 2] = 1.0
x_min, x_max, y_min, y_max = self.config['pred_range']
if feat[-1, 0] < x_min or feat[-1, 0] > x_max or feat[
-1, 1] < y_min or feat[-1, 1] > y_max:
continue
ctrs.append(feat[-1, :2].copy())
feat[1:, :2] -= feat[:-1, :2]
feat[step[0], :2] = 0
feats.append(feat)
gt_preds.append(gt_pred)
has_preds.append(has_pred)
feats = np.asarray(feats, np.float32)
ctrs = np.asarray(ctrs, np.float32)
gt_preds = np.asarray(gt_preds, np.float32)
has_preds = np.asarray(has_preds, np.bool)
data['feats'] = feats
data['ctrs'] = ctrs
data['orig'] = orig
data['theta'] = theta
data['rot'] = rot
data['gt_preds'] = gt_preds
data['has_preds'] = has_preds
return data
def get_lane_graph(self, data):
"""Get a rectangle area defined by pred_range."""
x_min, x_max, y_min, y_max = self.config['pred_range']
radius = max(abs(x_min), abs(x_max)) + max(abs(y_min), abs(y_max))
lane_ids = self.am.get_lane_ids_in_xy_bbox(data['orig'][0],
data['orig'][1],
data['city'], radius)
lane_ids = copy.deepcopy(lane_ids)
lanes = dict()
for lane_id in lane_ids:
lane = self.am.city_lane_centerlines_dict[data['city']][lane_id]
lane = copy.deepcopy(lane)
centerline = np.matmul(data['rot'],
(lane.centerline -
data['orig'].reshape(-1, 2)).T).T
x, y = centerline[:, 0], centerline[:, 1]
if x.max() < x_min or x.min() > x_max or y.max() < y_min or y.min(
) > y_max:
continue
else:
"""Getting polygons requires original centerline"""
polygon = self.am.get_lane_segment_polygon(
lane_id, data['city'])
polygon = copy.deepcopy(polygon)
lane.centerline = centerline
lane.polygon = np.matmul(data['rot'],
(polygon[:, :2] -
data['orig'].reshape(-1, 2)).T).T
lanes[lane_id] = lane
lane_ids = list(lanes.keys())
ctrs, feats, turn, control, intersect = [], [], [], [], []
for lane_id in lane_ids:
lane = lanes[lane_id]
ctrln = lane.centerline
num_segs = len(ctrln) - 1
ctrs.append(np.asarray((ctrln[:-1] + ctrln[1:]) / 2.0, np.float32))
feats.append(np.asarray(ctrln[1:] - ctrln[:-1], np.float32))
x = np.zeros((num_segs, 2), np.float32)
if lane.turn_direction == 'LEFT':
x[:, 0] = 1
elif lane.turn_direction == 'RIGHT':
x[:, 1] = 1
else:
pass
turn.append(x)
control.append(lane.has_traffic_control *
np.ones(num_segs, np.float32))
intersect.append(lane.is_intersection *
np.ones(num_segs, np.float32))
node_idcs = []
count = 0
for i, ctr in enumerate(ctrs):
node_idcs.append(range(count, count + len(ctr)))
count += len(ctr)
num_nodes = count
pre, suc = dict(), dict()
for key in ['u', 'v']:
pre[key], suc[key] = [], []
for i, lane_id in enumerate(lane_ids):
lane = lanes[lane_id]
idcs = node_idcs[i]
pre['u'] += idcs[1:]
pre['v'] += idcs[:-1]
if lane.predecessors is not None:
for nbr_id in lane.predecessors:
if nbr_id in lane_ids:
j = lane_ids.index(nbr_id)
pre['u'].append(idcs[0])
pre['v'].append(node_idcs[j][-1])
suc['u'] += idcs[:-1]
suc['v'] += idcs[1:]
if lane.successors is not None:
for nbr_id in lane.successors:
if nbr_id in lane_ids:
j = lane_ids.index(nbr_id)
suc['u'].append(idcs[-1])
suc['v'].append(node_idcs[j][0])
lane_idcs = []
for i, idcs in enumerate(node_idcs):
lane_idcs.append(i * np.ones(len(idcs), np.int64))
lane_idcs = np.concatenate(lane_idcs, 0)
pre_pairs, suc_pairs, left_pairs, right_pairs = [], [], [], []
for i, lane_id in enumerate(lane_ids):
lane = lanes[lane_id]
nbr_ids = lane.predecessors
if nbr_ids is not None:
for nbr_id in nbr_ids:
if nbr_id in lane_ids:
j = lane_ids.index(nbr_id)
pre_pairs.append([i, j])
nbr_ids = lane.successors
if nbr_ids is not None:
for nbr_id in nbr_ids:
if nbr_id in lane_ids:
j = lane_ids.index(nbr_id)
suc_pairs.append([i, j])
nbr_id = lane.l_neighbor_id
if nbr_id is not None:
if nbr_id in lane_ids:
j = lane_ids.index(nbr_id)
left_pairs.append([i, j])
nbr_id = lane.r_neighbor_id
if nbr_id is not None:
if nbr_id in lane_ids:
j = lane_ids.index(nbr_id)
right_pairs.append([i, j])
pre_pairs = np.asarray(pre_pairs, np.int64)
suc_pairs = np.asarray(suc_pairs, np.int64)
left_pairs = np.asarray(left_pairs, np.int64)
right_pairs = np.asarray(right_pairs, np.int64)
graph = dict()
graph['ctrs'] = np.concatenate(ctrs, 0)
graph['num_nodes'] = num_nodes
graph['feats'] = np.concatenate(feats, 0)
graph['turn'] = np.concatenate(turn, 0)
graph['control'] = np.concatenate(control, 0)
graph['intersect'] = np.concatenate(intersect, 0)
graph['pre'] = [pre]
graph['suc'] = [suc]
graph['lane_idcs'] = lane_idcs
graph['pre_pairs'] = pre_pairs
graph['suc_pairs'] = suc_pairs
graph['left_pairs'] = left_pairs
graph['right_pairs'] = right_pairs
for k1 in ['pre', 'suc']:
for k2 in ['u', 'v']:
graph[k1][0][k2] = np.asarray(graph[k1][0][k2], np.int64)
for key in ['pre', 'suc']:
if 'scales' in self.config and self.config['scales']:
#TODO: delete here
graph[key] += dilated_nbrs2(graph[key][0], graph['num_nodes'],
self.config['scales'])
else:
graph[key] += dilated_nbrs(graph[key][0], graph['num_nodes'],
self.config['num_scales'])
return graph
def verify_halluc_lane_extent_index(enable_lane_boundaries=False):
"""
"""
avm = ArgoverseMap()
city_names = ["MIA", "PIT"]
for city_name in city_names:
# get all lane segment IDs inside of this city
lane_segment_ids = list(
avm.city_lane_centerlines_dict[city_name].keys())
for lane_segment_id in lane_segment_ids:
xmin, ymin, xmax, ymax = find_lane_segment_bounds_in_table(
adm, city_name, lane_segment_id)
predecessor_ids = avm.get_lane_segment_predecessor_ids(
lane_segment_id, city_name)
successor_ids = avm.get_lane_segment_successor_ids(
lane_segment_id, city_name)
(r_neighbor_id, l_neighbor_id) = avm.get_lane_segment_adjacent_ids(
lane_segment_id, city_name)
lane_centerline = avm.get_lane_segment_centerline(
lane_segment_id, city_name)
halluc_lane_polygon = avm.get_lane_segment_polygon(
lane_segment_id, city_name)
fig = plt.figure(figsize=(22.5, 8))
ax = fig.add_subplot(111)
# add the lane of interest
add_lane_segment_to_ax(ax, lane_centerline, halluc_lane_polygon,
"y", xmin, xmax, ymin, ymax)
if predecessor_ids is not None:
# add predecessors
for predecessor_id in predecessor_ids:
lane_centerline = avm.get_lane_segment_centerline(
predecessor_id, city_name)
halluc_lane_polygon = avm.get_lane_segment_polygon(
predecessor_id, city_name)
xmin, ymin, xmax, ymax = find_lane_segment_bounds_in_table(
adm, city_name, predecessor_id)
add_lane_segment_to_ax(ax, lane_centerline,
halluc_lane_polygon, "r", xmin,
xmax, ymin, ymax)
if successor_ids is not None:
# add successors
for successor_id in successor_ids:
lane_centerline = avm.get_lane_segment_centerline(
successor_id, city_name)
halluc_lane_polygon = avm.get_lane_segment_polygon(
successor_id, city_name)
xmin, ymin, xmax, ymax = find_lane_segment_bounds_in_table(
adm, city_name, successor_id)
add_lane_segment_to_ax(ax, lane_centerline,
halluc_lane_polygon, "b", xmin,
xmax, ymin, ymax)
# add left neighbor
if l_neighbor_id is not None:
lane_centerline = avm.get_lane_segment_centerline(
l_neighbor_id, city_name)
halluc_lane_polygon = avm.get_lane_segment_polygon(
l_neighbor_id, city_name)
xmin, ymin, xmax, ymax = find_lane_segment_bounds_in_table(
adm, city_name, l_neighbor_id)
add_lane_segment_to_ax(ax, lane_centerline,
halluc_lane_polygon, "g", xmin, xmax,
ymin, ymax)
# add right neighbor
if r_neighbor_id is not None:
lane_centerline = avm.get_lane_segment_centerline(
r_neighbor_id, city_name)
halluc_lane_polygon = avm.get_lane_segment_polygon(
r_neighbor_id, city_name)
xmin, ymin, xmax, ymax = find_lane_segment_bounds_in_table(
adm, city_name, r_neighbor_id)
add_lane_segment_to_ax(ax, lane_centerline,
halluc_lane_polygon, "m", xmin, xmax,
ymin, ymax)
if enable_lane_boundaries:
# Compare with Argo's proprietary, ground truth lane boundaries
gt_lane_polygons = avm.city_to_lane_polygons_dict[city_name]
for gt_lane_polygon in gt_lane_polygons:
dist = np.linalg.norm(
gt_lane_polygon.mean(axis=0)[:2] -
np.array([xmin, ymin]))
if dist < 30:
ax.plot(gt_lane_polygon[:, 0],
gt_lane_polygon[:, 1],
color="k",
alpha=0.3,
zorder=1)
ax.axis("equal")
plt.show()
datetime_str = generate_datetime_string()
plt.savefig(
f"lane_segment_id_{lane_segment_id}_@_{datetime_str}.jpg")
plt.close("all")
class GraphExtractor(object):
def __init__(self, config, mode='train'):
self.am = ArgoverseMap()
self.config = config
self.mode = mode
def __del__(self):
del self.am
def dilated_nbrs(self, nbr, num_nodes, num_scales):
data = np.ones(len(nbr['u']), np.bool)
csr = sparse.csr_matrix((data, (nbr['u'], nbr['v'])),
shape=(num_nodes, num_nodes))
mat = csr
nbrs = []
for i in range(1, num_scales):
mat = mat * mat
nbr = dict()
coo = mat.tocoo()
nbr['u'] = coo.row.astype(np.int16)
nbr['v'] = coo.col.astype(np.int16)
nbrs.append(nbr)
return nbrs
def extract(self, data):
"""Get a rectangle area defined by pred_range."""
x_min, x_max, y_min, y_max = self.config['pred_range']
radius = max(abs(x_min), abs(x_max)) + max(abs(y_min), abs(y_max))
lane_ids = self.am.get_lane_ids_in_xy_bbox(data['orig'][0],
data['orig'][1],
data['city'], radius)
lane_ids = copy.deepcopy(lane_ids)
"""Get all lane within self.config['pred_range'], convert centerline and polygon to rotated and biased"""
# what's polygon
lanes = dict()
for lane_id in lane_ids:
lane = self.am.city_lane_centerlines_dict[data['city']][lane_id]
lane = copy.deepcopy(lane)
centerline = np.matmul(data['rot'],
(lane.centerline -
data['orig'].reshape(-1, 2)).T).T
x, y = centerline[:, 0], centerline[:, 1]
if x.max() < x_min or x.min() > x_max or y.max() < y_min or y.min(
) > y_max:
continue
else:
"""Getting polygons requires original centerline"""
polygon = self.am.get_lane_segment_polygon(
lane_id, data['city'])
polygon = copy.deepcopy(polygon)
lane.centerline = centerline
lane.polygon = np.matmul(data['rot'],
(polygon[:, :2] -
data['orig'].reshape(-1, 2)).T).T
lanes[lane_id] = lane
"""Lane feature: ctrs(position), feats(shape), turn, control, intersect"""
lane_ids = list(lanes.keys())
ctrs, feats, turn, control, intersect = [], [], [], [], []
for lane_id in lane_ids:
lane = lanes[lane_id]
ctrln = lane.centerline
num_segs = len(ctrln) - 1
ctrs.append(np.asarray((ctrln[:-1] + ctrln[1:]) / 2.0, np.float32))
feats.append(np.asarray(ctrln[1:] - ctrln[:-1], np.float32))
x = np.zeros((num_segs, 2), np.float32)
if lane.turn_direction == 'LEFT':
x[:, 0] = 1
elif lane.turn_direction == 'RIGHT':
x[:, 1] = 1
else:
pass
turn.append(x)
control.append(lane.has_traffic_control *
np.ones(num_segs, np.float32))
intersect.append(lane.is_intersection *
np.ones(num_segs, np.float32))
# -------------------------node_idcs---------------------
node_idcs = []
count = 0
for ctr in ctrs: # node_idcs: list, i-th element: i-th lane nodes ids
node_idcs.append(range(count, count + len(ctr)))
count += len(ctr)
num_nodes = count
# -------------------------lane_idcs---------------------
# lane[idc] = a means idc-th node belongs to the a-th lane
lane_idcs = []
for i, idcs in enumerate(node_idcs):
lane_idcs.append(
i *
np.ones(len(idcs), np.int64)) # TODO: what does lane_idcs do?
lane_idcs = np.concatenate(lane_idcs, 0)
# **********************************Map Related work***************************
# =========================================
# ==============Hdmap Graph Build==========
# =========================================
# -------all in all, pairs is for lanes; no pairs is for lanes--------
# ---------------------------pre and suc for lanes--------------------
pre_pairs, suc_pairs, left_pairs, right_pairs = [], [], [], []
for i, lane_id in enumerate(lane_ids):
lane = lanes[lane_id]
nbr_ids = lane.predecessors
if nbr_ids is not None:
for nbr_id in nbr_ids:
if nbr_id in lane_ids:
j = lane_ids.index(nbr_id)
pre_pairs.append([i, j])
nbr_ids = lane.successors
if nbr_ids is not None:
for nbr_id in nbr_ids:
if nbr_id in lane_ids:
j = lane_ids.index(nbr_id)
suc_pairs.append([i, j])
nbr_id = lane.l_neighbor_id
if nbr_id is not None:
if nbr_id in lane_ids:
j = lane_ids.index(nbr_id)
left_pairs.append([i, j])
nbr_id = lane.r_neighbor_id
if nbr_id is not None:
if nbr_id in lane_ids:
j = lane_ids.index(nbr_id)
right_pairs.append([i, j])
pre_pairs = np.asarray(pre_pairs, np.int16)
suc_pairs = np.asarray(suc_pairs, np.int16)
left_pairs = np.asarray(left_pairs, np.int16)
right_pairs = np.asarray(right_pairs, np.int16)
# ---------------------------pre and suc for nodes--------------------
pre, suc = dict(), dict()
for key in ['u', 'v']:
pre[key], suc[key] = [], []
for i, lane_id in enumerate(lane_ids):
lane = lanes[lane_id]
idcs = node_idcs[i]
pre['u'] += idcs[1:]
pre['v'] += idcs[:-1]
if lane.predecessors is not None:
for nbr_id in lane.predecessors:
if nbr_id in lane_ids:
j = lane_ids.index(nbr_id)
pre['u'].append(idcs[0])
pre['v'].append(
node_idcs[j]
[-1]) # v is the pre of u, v is src, u is dest
suc['u'] += idcs[:-1]
suc['v'] += idcs[1:]
if lane.successors is not None:
for nbr_id in lane.successors:
if nbr_id in lane_ids:
j = lane_ids.index(nbr_id)
suc['u'].append(idcs[-1])
suc['v'].append(node_idcs[j][0])
pre['u'] = np.asarray(pre['u'], dtype=np.int16)
pre['v'] = np.asarray(pre['v'], dtype=np.int16)
suc['u'] = np.asarray(suc['u'], dtype=np.int16)
suc['v'] = np.asarray(suc['v'], dtype=np.int16)
# -------------------dilate pre and suc: opition 1--------------------
dilated_pre = [pre]
dilated_pre += self.dilated_nbrs(pre, num_nodes,
self.config['num_scales'])
dilated_suc = [suc]
dilated_suc += self.dilated_nbrs(suc, num_nodes,
self.config['num_scales'])
# --------------------build nodes left and right graph-----------------
num_lanes = lane_idcs[-1].item() + 1
left, right = dict(), dict()
dist = np.expand_dims(ctrs, axis=1) - np.expand_dims(ctrs, axis=0)
dist = np.sqrt((dist**2).sum(2))
hi = np.arange(num_nodes).reshape(-1, 1).repeat(num_nodes,
axis=1).reshape(-1)
wi = np.arange(num_nodes).reshape(1, -1).repeat(num_nodes,
axis=0).reshape(-1)
row_idcs = np.arange(num_nodes)
pre_mat = np.zeros((num_lanes, num_lanes))
pre_mat[pre_pairs[:, 0], pre_pairs[:, 1]] = 1
suc_mat = np.zeros((num_lanes, num_lanes))
suc_mat[suc_pairs[:, 0], suc_pairs[:, 1]] = 1
pairs = left_pairs
if len(pairs) > 0:
# construct lane left graph
mat = np.zeros((num_lanes, num_lanes))
mat[pairs[:, 0], pairs[:, 1]] = 1
mat = (
np.matmul(mat, pre_mat) + np.matmul(mat, suc_mat) + mat
) > 0.5 # left lane's suc or pre lane is also self's left lane
# filter with distance
left_dist = dist.copy()
# if lane j is the lane i's left, then all nodes in lane j is the left of any node in lane i
mask = np.logical_not(mat[lane_idcs[hi], lane_idcs[wi]])
# set the distance between nodes that has no left relation are very vert large
left_dist[hi[mask], wi[mask]] = 1e6
# find the each node's nearest node
min_dist, min_idcs = left_dist.min(1), left_dist.argmin(1)
# if nearest node's distance > self.config['cross_dist'], then this node does not have left node
mask = min_dist < self.config['cross_dist']
# if the angle between nearest node is too big , the this node does not have left node
ui = row_idcs[mask]
vi = min_idcs[mask]
f1 = feats[ui]
f2 = feats[vi]
t1 = np.arctan2(f1[:, 1], f1[:, 0])
t2 = np.arctan2(f2[:, 1], f2[:, 0])
dt = np.abs(t1 - t2)
m = dt > np.pi
dt[m] = np.abs(dt[m] - 2 * np.pi)
m = dt < 0.25 * np.pi
ui = ui[m]
vi = vi[m]
left['u'] = ui.astype(
np.int16) # u is the idx of node that has left neighbor
left['v'] = vi.astype(
np.int16) # v[i] is the idx of left neighbor of node u[i]
else:
left['u'] = np.zeros(0, np.int16)
left['v'] = np.zeros(0, np.int16)
pairs = right_pairs
if len(pairs) > 0:
mat = np.zeros((num_lanes, num_lanes))
mat[pairs[:, 0], pairs[:, 1]] = 1
mat = (np.matmul(mat, pre_mat) + np.matmul(mat, suc_mat) +
mat) > 0.5
right_dist = dist.copy()
mask = np.logical_not(mat[lane_idcs[hi], lane_idcs[wi]])
right_dist[hi[mask], wi[mask]] = 1e6
min_dist, min_idcs = right_dist.min(1), right_dist.argmin(1)
mask = min_dist < self.config['cross_dist']
ui = row_idcs[mask]
vi = min_idcs[mask]
f1 = feats[ui]
f2 = feats[vi]
t1 = np.arctan2(f1[:, 1], f1[:, 0])
t2 = np.arctan2(f2[:, 1], f2[:, 0])
dt = np.abs(t1 - t2)
m = dt > np.pi
dt[m] = np.abs(dt[m] - 2 * np.pi)
m = dt < 0.25 * np.pi
ui = ui[m]
vi = vi[m]
right['u'] = ui.astype(np.int16)
right['v'] = vi.astype(np.int16)
else:
right['u'] = np.zeros(0, np.int16)
right['v'] = np.zeros(0, np.int16)
graph = dict()
graph['num_nodes'] = num_nodes
# map node feats
graph['ctrs'] = ctrs
graph['feats'] = feats
graph['turn'] = np.concatenate(turn, 0)
graph['control'] = np.concatenate(control, 0)
graph['intersect'] = np.concatenate(intersect, 0)
# map node graph
graph['pre'] = dilated_pre
graph['suc'] = dilated_suc
graph['left'] = left
graph['right'] = right
# lane pairs
graph['lane_idcs'] = lane_idcs
graph['pre_pairs'] = pre_pairs
graph['suc_pairs'] = suc_pairs
graph['left_pairs'] = left_pairs
graph['right_pairs'] = right_pairs
return graph
class ArgoverseForecastDataset(torch.utils.data.Dataset):
def __init__(self, cfg):
super().__init__()
self.am = ArgoverseMap()
self.axis_range = self.get_map_range(self.am)
self.city_halluc_bbox_table, self.city_halluc_tableidx_to_laneid_map = self.am.build_hallucinated_lane_bbox_index(
)
self.laneid_map = self.process_laneid_map()
self.vector_map, self.extra_map = self.generate_vector_map()
# am.draw_lane(city_halluc_tableidx_to_laneid_map['PIT']['494'], 'PIT')
# self.save_vector_map(self.vector_map)
self.last_observe = cfg['last_observe']
##set root_dir to the correct path to your dataset folder
self.root_dir = cfg['data_locate']
self.afl = ArgoverseForecastingLoader(self.root_dir)
self.map_feature = dict(PIT=[], MIA=[])
self.city_name, self.center_xy, self.rotate_matrix = dict(), dict(
), dict()
def __len__(self):
return len(self.afl)
def __getitem__(self, index):
# self.am.find_local_lane_polygons()
self.trajectory, city_name, extra_fields = self.get_trajectory(index)
traj_id = extra_fields['trajectory_id'][0]
self.city_name.update({str(traj_id): city_name})
center_xy = self.trajectory[self.last_observe - 1][1]
self.center_xy.update({str(traj_id): center_xy})
trajectory_feature = (self.trajectory -
np.array(center_xy).reshape(1, 1, 2)).reshape(
-1, 4)
rotate_matrix = get_rotate_matrix(
trajectory_feature[self.last_observe, :]) # rotate coordinate
self.rotate_matrix.update({str(traj_id): rotate_matrix})
trajectory_feature = ((trajectory_feature.reshape(-1, 2)).dot(
rotate_matrix.T)).reshape(-1, 4)
trajectory_feature = self.normalize_coordinate(
trajectory_feature, city_name) # normalize to [-1, 1]
self.traj_feature = torch.from_numpy(
np.hstack((
trajectory_feature,
extra_fields['TIMESTAMP'].reshape(-1, 1),
# extra_fields['OBJECT_TYPE'].reshape(-1, 1),
extra_fields['trajectory_id'].reshape(-1, 1)))).float()
map_feature_dict = dict(PIT=[], MIA=[])
for city in ['PIT', 'MIA']:
for i in range(len(self.vector_map[city])):
map_feature = (self.vector_map[city][i] -
np.array(center_xy).reshape(1, 1, 2)).reshape(
-1, 2)
map_feature = (map_feature.dot(rotate_matrix.T)).reshape(-1, 4)
map_feature = self.normalize_coordinate(map_feature, city)
tmp_tensor = torch.from_numpy(
np.hstack((
map_feature,
self.extra_map[city]['turn_direction'][i],
self.extra_map[city]['in_intersection'][i],
self.extra_map[city]['has_traffic_control'][i],
# self.extra_map[city]['OBJECT_TYPE'][i],
self.extra_map[city]['lane_id'][i])))
map_feature_dict[city].append(tmp_tensor.float())
# self.map_feature[city] = np.array(self.map_feature[city])
self.map_feature[city] = map_feature_dict[city]
self.map_feature['city_name'] = city_name
return self.traj_feature, self.map_feature
def get_trajectory(self, index):
seq_path = self.afl.seq_list[index]
data = self.afl.get(seq_path).seq_df
data = data[data['OBJECT_TYPE'] == 'AGENT']
extra_fields = dict(TIMESTAMP=[], OBJECT_TYPE=[], trajectory_id=[])
polyline = []
j = int(str(seq_path).split('/')[-1].split('.')[0])
flag = True
city_name = ''
for _, row in data.iterrows():
if flag:
xlast = row['X']
ylast = row['Y']
tlast = row['TIMESTAMP']
city_name = row['CITY_NAME']
flag = False
continue
startpoint = np.array([xlast, ylast])
endpoint = np.array([row['X'], row['Y']])
xlast = row['X']
ylast = row['Y']
extra_fields['TIMESTAMP'].append(tlast)
extra_fields['OBJECT_TYPE'].append(0) # 'AGENT'
extra_fields['trajectory_id'].append(j) # 'AGENT'
tlast = row['TIMESTAMP']
polyline.append([startpoint, endpoint])
extra_fields['TIMESTAMP'] = np.array(extra_fields['TIMESTAMP'])
extra_fields['TIMESTAMP'] -= np.min(
extra_fields['TIMESTAMP']) # adjust time stamp
extra_fields['OBJECT_TYPE'] = np.array(extra_fields['OBJECT_TYPE'])
extra_fields['trajectory_id'] = np.array(extra_fields['trajectory_id'])
return np.array(polyline), city_name, extra_fields
def generate_vector_map(self):
vector_map = {'PIT': [], 'MIA': []}
extra_map = {
'PIT':
dict(OBJECT_TYPE=[],
turn_direction=[],
lane_id=[],
in_intersection=[],
has_traffic_control=[]),
'MIA':
dict(OBJECT_TYPE=[],
turn_direction=[],
lane_id=[],
in_intersection=[],
has_traffic_control=[])
}
polyline = []
# index = 1
pbar = tqdm(total=17326)
pbar.set_description("Generating Vector Map")
for city_name in ['PIT', 'MIA']:
for key in self.laneid_map[city_name]:
pts = self.am.get_lane_segment_polygon(key, city_name)
turn_str = self.am.get_lane_turn_direction(key, city_name)
if turn_str == 'LEFT':
turn = -1
elif turn_str == 'RIGHT':
turn = 1
else:
turn = 0
pts_len = pts.shape[0] // 2
positive_pts = pts[:pts_len, :2]
negative_pts = pts[pts_len:2 * pts_len, :2]
polyline.clear()
for i in range(pts_len - 1):
v1 = np.array([positive_pts[i], positive_pts[i + 1]])
v2 = np.array([
negative_pts[pts_len - 1 - i],
negative_pts[pts_len - i - 2]
])
polyline.append(v1)
polyline.append(v2)
# extra_field['table_index'] = self.laneid_map[city_name][key]
repeat_t = 2 * (pts_len - 1)
vector_map[city_name].append(np.array(polyline).copy())
extra_map[city_name]['turn_direction'].append(
np.repeat(turn, repeat_t, axis=0).reshape(-1, 1))
extra_map[city_name]['OBJECT_TYPE'].append(
np.repeat(-1, repeat_t, axis=0).reshape(-1, 1)) #HD Map
extra_map[city_name]['lane_id'].append(
np.repeat(int(key), repeat_t, axis=0).reshape(-1, 1))
extra_map[city_name]['in_intersection'].append(
np.repeat(1 *
self.am.lane_is_in_intersection(key, city_name),
repeat_t,
axis=0).reshape(-1, 1))
extra_map[city_name]['has_traffic_control'].append(
np.repeat(1 * self.am.lane_has_traffic_control_measure(
key, city_name),
repeat_t,
axis=0).reshape(-1, 1))
# if index > 10:
# break
# index = index + 1
pbar.update(1)
pbar.close()
print("Generate Vector Map Successfully!")
return vector_map, extra_map #vector_map:list
def process_laneid_map(self):
laneid_map = {}
tmp_map = {}
tmp1_map = {}
for key in self.city_halluc_tableidx_to_laneid_map['PIT']:
tmp_map[self.city_halluc_tableidx_to_laneid_map['PIT'][key]] = key
laneid_map['PIT'] = tmp_map
for key in self.city_halluc_tableidx_to_laneid_map['MIA']:
tmp1_map[self.city_halluc_tableidx_to_laneid_map['MIA'][key]] = key
laneid_map['MIA'] = tmp1_map
return laneid_map
def get_map_range(self, am):
map_range = dict(PIT={}, MIA={})
for city_name in ['PIT', 'MIA']:
poly = am.get_vector_map_lane_polygons(city_name)
poly_modified = (np.vstack(poly))[:, :2]
max_coordinate = np.max(poly_modified, axis=0)
min_coordinate = np.min(poly_modified, axis=0)
map_range[city_name].update({'max': max_coordinate})
map_range[city_name].update({'min': min_coordinate})
return map_range
def normalize_coordinate(self, array, city_name):
max_coordinate = self.axis_range[city_name]['max']
min_coordinate = self.axis_range[city_name]['min']
array = (10. * (array.reshape(-1, 2)) /
(max_coordinate - min_coordinate)).reshape(-1, 4)
return array
def save_vector_map(self, vector_map):
save_path = "./data/vector_map/"
for city_name in ['PIT', 'MIA']:
tmp_map = np.vstack(vector_map[city_name]).reshape(-1, 4)
np.save(save_path + city_name + "_vectormap", tmp_map)
def main(data_dir):
"""
"""
fnames = glob.glob(f"{data_dir}/*.csv")
fnames = [Path(fname).name for fname in fnames]
am = ArgoverseMap()
city_graph_dict = build_city_lane_graphs(am)
for fname in fnames:
# # very hard cases
# if int(Path(fname).stem) not in [
# 166633, 150381,11905, 136010, 49854, 27155]:
# continue
# # hard cases -- ,
# [174545,119781, 210709, 139445, 11381, 175883, 122703, 166633]: #23333,,124414]:
#
csv_fpath = f"{data_dir}/{fname}"
traj, city_name = get_traj_and_city_name_from_csv(csv_fpath)
plausible_start_ids = set()
lane_vote_dict = defaultdict(int)
for j, pt in enumerate(traj):
contained_ids = am.get_lane_segments_containing_xy(
pt[0], pt[1], city_name)
for id in contained_ids:
lane_vote_dict[id] += 1
plausible_start_ids.add(id)
plausible_start_ids = list(plausible_start_ids)
plausible_start_ids.sort()
paths = []
# START BFS IN ANY PLAUSIBLE LANE ID!
for start_id in plausible_start_ids:
paths.extend(
find_all_paths_from_src(city_graph_dict[city_name],
str(start_id),
max_depth=DFS_MAX_DEPTH))
paths = convert_str_lists_to_int_lists(paths)
paths = trim_paths_with_no_inliers(paths, lane_vote_dict)
paths = remove_repeated_paths(paths)
path_votes_dict = defaultdict(int)
for path_id, path in enumerate(paths):
for id in path:
path_votes_dict[path_id] += lane_vote_dict[id]
# find which path has most inliers
best_path_ids = get_dict_key_with_max_value(path_votes_dict)
# if they are all tied, take the shortest
best_path_lengths = [len(paths[id]) for id in best_path_ids]
min_best_path_length = min(best_path_lengths)
best_path_ids = [
id for id in best_path_ids
if len(paths[id]) == min_best_path_length
]
fig = plt.figure(figsize=(15, 15))
plt.axis("off")
ax = fig.add_subplot(111)
plot_all_nearby_lanes(am, ax, city_name, np.mean(traj[:, 0]),
np.mean(traj[:, 1]))
colors = ["g", "b", "r", "m"]
# then plot this path
for best_path_id in best_path_ids:
color = colors[best_path_id % len(colors)]
print(
"Candidate: ",
paths[best_path_id],
" with ",
path_votes_dict[best_path_id],
)
for lane_id in paths[best_path_id]:
polygon_3d = am.get_lane_segment_polygon(lane_id, city_name)
plot_lane_segment_patch(polygon_3d[:, :2], ax, color=color)
ax.text(np.mean(polygon_3d[:, 0]), np.mean(polygon_3d[:, 1]),
f"{lane_id}")
# just use one for now
break
# draw_lane_ids(plausible_start_ids, am, ax, city_name)
all_nearby_lane_ids = []
for path in paths:
all_nearby_lane_ids.extend(path)
draw_lane_ids(set(all_nearby_lane_ids), am, ax, city_name)
draw_traj(traj, ax)
xmin, ymin, xmax, ymax = compute_point_cloud_bbox(traj)
WINDOW_RADIUS_MARGIN = 10
xmin -= WINDOW_RADIUS_MARGIN
xmax += WINDOW_RADIUS_MARGIN
ymin -= WINDOW_RADIUS_MARGIN
ymax += WINDOW_RADIUS_MARGIN
ax.set_xlim([xmin, xmax])
ax.set_ylim([ymin, ymax])
plt.savefig(
f"/Users/johnlamb/Documents/argoverse-api/temp_files_oracle/{Path(fname).stem}.png"
)
plt.close("all")
]
]
# lines = [ headstart+build_line(line, lane_centerlines_dict, reverse=(not front)) for line in lines]
# # Label build
# label_line = crop_line( lines[label], source[-2], target[-1])
# label_line_distance = line_distance(label_line) + 5
# #endregion
# # Distance threshold
# lines = [ line_distance_crop(line, label_line_distance) for line in lines]
#region Build intersection polygon
intersection = Polygon(
avm.get_lane_segment_polygon(intersection_ids[0], city)[:, 0:2])
for id in intersection_ids[1:]:
intersection = intersection.intersection(
Polygon(avm.get_lane_segment_polygon(id, city)[:, 0:2]))
polygons = [intersection.exterior.coords]
#endregion
#region Find Separation Point
separation_point = find_separation_point([
build_line([id], lane_centerlines_dict, reverse=(not front))
for id in intersection_ids
])
#endregion