def __getitem__(self, index):
current_loader = self.afl.get(self.seq_paths[index])
agent_traj = current_loader.agent_traj
candidate_centerlines = self.avm.get_candidate_centerlines_for_traj(
agent_traj, current_loader.city, viz=False)
if self.oracle:
candidate_centerlines = [
get_oracle_from_candidate_centerlines(candidate_centerlines,
agent_traj)
]
if self.mode_test:
seq_index = int(
os.path.basename(self.seq_paths[index]).split('.')[0])
agent_train_traj = agent_traj[:self.train_seq_size, :]
all_centerline_traj = []
for centerline in candidate_centerlines:
all_centerline_traj.append(
torch.Tensor(
get_nt_distance(agent_train_traj,
current_centerline)).float())
return {
'seq_index': seq_index,
'train_agent': all_centerline_traj,
'centerline': candidate_centerlines,
'city': current_loader.city
}
else:
agent_train_traj = agent_traj[:self.train_seq_size, :]
agent_gt_traj = agent_traj[self.train_seq_size:, ]
all_centerline_train_traj = []
all_centerline_gt_traj = []
for centerline in candidate_centerlines:
all_centerline_train_traj.append(
torch.Tensor(
get_nt_distance(agent_train_traj,
current_centerline)).float())
all_centerline_gt_traj.append(
torch.Tensor(
get_nt_distance(agent_gt_traj,
current_centerline)).float())
agent_unnorm_gt_traj = torch.Tensor(
agent_traj[self.train_seq_size:, ]).float()
return {
'train_agent': all_centerline_train_traj,
'gt_agent': all_centerline_gt_traj,
'gt_unnorm_agent': agent_unnorm_gt_traj,
'centerline': current_centerline,
'city': current_loader.city
}
def compute_map_features(
agent_track: np.ndarray,
oracle_centerline: np.ndarray,
obs_len: int,
pred_len: int,
raw_data_format: Dict[str, int],
):
agent_obs = agent_track[:obs_len]
agent_xy_obs = agent_obs[:, [raw_data_format["X"], raw_data_format["Y"]
]].astype("float")
oracle_nt_dist = get_nt_distance(agent_xy_obs,
oracle_centerline,
viz=False)
oracle_nt_dist_norm = oracle_nt_dist - oracle_nt_dist[0, :]
delta_ref = copy.deepcopy(oracle_nt_dist[0, :])
for i in range(agent_xy_obs.shape[0] - 1, 0, -1):
oracle_nt_dist[i, :] = oracle_nt_dist[i, :] - oracle_nt_dist[i - 1, :]
oracle_nt_dist[0, :] = 0
angle_w_cl = np.zeros((agent_xy_obs.shape[0], 1))
angle_w_cl[1:, 0] = np.arctan2(oracle_nt_dist[1:, 1], oracle_nt_dist[1:,
0])
angle_w_cl[0, :] = angle_w_cl[1, :]
# angle_w_cl[np.isnan(angle_w_cl)] = np.pi/2
map_features = np.concatenate((oracle_nt_dist_norm, angle_w_cl), axis=1)
return map_features, delta_ref
def test_get_nt_distance():
"""Compute distances in centerline frame for a trajectory"""
"""Test Case
0 1 . . 3 . 4 . 5 . 6 . 7
* 5
\
x * 4
\
* x 3
\
x * x 2
\
. *---*---*---*---x 1
x 0
"""
xy = np.array([(0.0, 4.0), (3.0, 3.0), (1.0, 2.0), (4.0, 2.0), (5.0, 0.0),
(7.0, 1.0)])
centerline = np.array([(1.0, 5.0), (1.5, 4.0), (2.0, 3.0), (2.5, 2.0),
(3.0, 1.0), (4.0, 1.0), (5.0, 1.0), (6.0, 1.0),
(7.0, 1.0)])
nt_dist = get_nt_distance(xy, centerline)
expected_nt_dist = np.array([[1.34, 0.44], [2.68, 0.89], [2.68, 1.34],
[5.47, 1.0], [6.47, 1.0], [8.47, 0.0]])
print(nt_dist, expected_nt_dist)
np.array_equal(nt_dist, expected_nt_dist)
def compute_features_old(self,
seq_path,
map_instance,
social_feature_instance,
avm,
mode="train"):
check1 = True
if check1:
if mode == "train" or mode == "validate":
current_loader = self.afl.get(seq_path)
agent_traj = current_loader.agent_traj
# df = pd.read_csv(seq_path, dtype={"TIMESTAMP": str})
# agent_track = df[df["OBJECT_TYPE"] == "AGENT"].values
candidate_centerlines = self.avm.get_candidate_centerlines_for_traj(
agent_traj, current_loader.city, viz=False)
current_centerline = get_oracle_from_candidate_centerlines(
candidate_centerlines, agent_traj)
agent_traj_norm = get_nt_distance(agent_traj,
current_centerline)
return None, agent_traj_norm, {
"ORACLE_CENTERLINE": current_centerline
}
elif mode == "validate_multiple":
current_loader = self.afl.get(seq_path)
agent_traj = current_loader.agent_traj
candidate_centerlines = self.avm.get_candidate_centerlines_for_traj(
agent_traj, current_loader.city, viz=False)
else:
map_features, map_feature_helpers = self.map_features_utils_instance.compute_map_features(
agent_track, 20, 50, RAW_DATA_FORMAT, mode, avm)
return None, map_features, map_feature_helpers
def __getitem__(self, index):
current_loader = self.afl.get(self.seq_paths[index])
agent_traj = current_loader.agent_traj
candidate_centerlines = self.avm.get_candidate_centerlines_for_traj(
agent_traj, current_loader.city, viz=False)
# if self.oracle:
current_centerline = get_oracle_from_candidate_centerlines(
candidate_centerlines, agent_traj)
# else:
# current_centerline=candidate_centerlines
if self.mode_test:
seq_index = int(
os.path.basename(self.seq_paths[index]).split('.')[0])
agent_train_traj = agent_traj[:self.train_seq_size, :]
agent_train_traj = get_nt_distance(agent_train_traj,
current_centerline)
agent_train_traj = torch.Tensor(agent_train_traj).float()
# gt_agent=self.get_coordinate_from_centerline(oracle_centerline,agent_train_traj)
return {
'seq_index': seq_index,
'train_agent': agent_train_traj,
'centerline': current_centerline,
'city': current_loader.city
}
else:
agent_train_traj = agent_traj[:self.train_seq_size, :]
agent_train_traj = get_nt_distance(agent_train_traj,
current_centerline)
agent_train_traj = torch.Tensor(agent_train_traj).float()
agent_gt_traj = agent_traj[self.train_seq_size:, ]
agent_gt_traj = get_nt_distance(agent_gt_traj, current_centerline)
agent_gt_traj = torch.Tensor(agent_gt_traj).float()
agent_unnorm_gt_traj = torch.Tensor(
agent_traj[self.train_seq_size:, ]).float()
return {
'seq_path': self.seq_paths[index],
'train_agent': agent_train_traj,
'gt_agent': agent_gt_traj,
'gt_unnorm_agent': agent_unnorm_gt_traj,
'centerline': current_centerline,
'city': current_loader.city
}
def get_centerline_nt_attributes(agent_traj, city, avm):
candidate_centerlines = avm.get_candidate_centerlines_for_traj(
agent_traj, city)
nt_distances = [
get_nt_distance(agent_traj, centerline)
for centerline in candidate_centerlines
]
return (candidate_centerlines, nt_distances)
def get_distances(start_coor, end_coor, centerlines):
#print("Start: ",start_coor)
#print("End: ",end_coor)
dist = get_nt_distance(np.array([start_coor, end_coor]),
np.array(centerlines))[1][1]
#print("Dist: ",dist)
return dist
def convert_neighbour_centerline(self, neighbours_traj, centerline, ref_t):
neighbour_centerline_frame = []
if len(neighbours_traj) == 0:
return np.array([])
for neighbour_traj in neighbours_traj:
temp = get_nt_distance(neighbour_traj, centerline)
temp[:, 1] = temp[:, 1] - ref_t
neighbour_centerline_frame.append(temp)
return np.stack(neighbour_centerline_frame, axis=0)
def compute_map_features(agent_track: np.ndarray, obs_len: int, pred_len: int,
raw_data_format: Dict[str, int], mode: str):
agent_xy = agent_track[:, [raw_data_format["X"], raw_data_format["Y"]
]].astype("float")
agent_obs = agent_track[:obs_len]
agent_xy_obs = agent_obs[:, [raw_data_format["X"], raw_data_format["Y"]
]].astype("float")
if mode == "test":
xy = agent_xy_obs
else:
xy = agent_xy
avm = ArgoverseMap()
city_name = agent_track[0, raw_data_format["CITY_NAME"]]
candidate_centerlines = avm.get_candidate_centerlines_for_traj(
xy,
city_name,
viz=False,
max_search_radius=_MAX_SEARCH_RADIUS_CENTERLINES,
)
oracle_centerline = get_oracle_from_candidate_centerlines(
candidate_centerlines, xy)
oracle_nt_dist = get_nt_distance(xy, oracle_centerline, viz=False)
oracle_nt_dist_norm = oracle_nt_dist - oracle_nt_dist[0, :]
delta_ref = copy.deepcopy(oracle_nt_dist[0, :])
for i in range(xy.shape[0] - 1, 0, -1):
oracle_nt_dist[i, :] = oracle_nt_dist[i, :] - oracle_nt_dist[i - 1, :]
oracle_nt_dist[0, :] = 0
angle_w_cl = np.zeros((xy.shape[0], 1))
angle_w_cl[1:, 0] = np.arctan2(oracle_nt_dist[1:, 1], oracle_nt_dist[1:,
0])
angle_w_cl[0, :] = angle_w_cl[1, :]
# angle_w_cl[np.isnan(angle_w_cl)] = np.pi/2
map_features = np.concatenate((oracle_nt_dist_norm, angle_w_cl), axis=1)
if mode == "test":
map_features = np.concatenate(
(map_features, np.full([pred_len, 3], None)), axis=0)
return map_features, oracle_centerline, delta_ref
def compute_map_features(
self,
agent_track: np.ndarray,
obs_len: int,
seq_len: int,
raw_data_format: Dict[str, int],
mode: str,
) -> Tuple[np.ndarray, Dict[str, Any]]:
"""Compute map based features for the given sequence.
If the mode is test, oracle_nt_dist will be empty, candidate_nt_dist will be populated.
If the mode is train/val, oracle_nt_dist will be populated, candidate_nt_dist will be empty.
Args:
agent_track : Data for the agent track
obs_len : Length of observed trajectory
seq_len : Length of the sequence
raw_data_format : Format of the sequence
mode: train/val/test mode
Returns:
oracle_nt_dist (numpy array): normal and tangential distances for oracle centerline
map_feature_helpers (dict): Dictionary containing helpers for map features
"""
# Get observed 2 secs of the agent
agent_xy = agent_track[:, [raw_data_format["X"], raw_data_format["Y"]
]].astype("float")
agent_track_obs = agent_track[:obs_len]
agent_xy_obs = agent_track_obs[:, [
raw_data_format["X"], raw_data_format["Y"]
]].astype("float")
# Get API for Argo Dataset map
avm = ArgoverseMap()
city_name = agent_track[0, raw_data_format["CITY_NAME"]]
# Get candidate centerlines using observed trajectory
if mode == "test":
oracle_centerline = np.full((seq_len, 2), None)
oracle_nt_dist = np.full((seq_len, 2), None)
candidate_centerlines = self.get_candidate_centerlines_for_trajectory(
agent_xy_obs,
city_name,
avm,
viz=False,
max_search_radius=self._MAX_SEARCH_RADIUS_CENTERLINES,
seq_len=seq_len,
max_candidates=self._MAX_CENTERLINE_CANDIDATES_TEST,
)
# Get nt distance for the entire trajectory using candidate centerlines
candidate_nt_distances = []
for candidate_centerline in candidate_centerlines:
candidate_nt_distance = np.full((seq_len, 2), None)
candidate_nt_distance[:obs_len] = get_nt_distance(
agent_xy_obs, candidate_centerline)
candidate_nt_distances.append(candidate_nt_distance)
else:
oracle_centerline = self.get_candidate_centerlines_for_trajectory(
agent_xy,
city_name,
avm,
viz=False,
max_search_radius=self._MAX_SEARCH_RADIUS_CENTERLINES,
seq_len=seq_len,
mode=mode,
)[0]
candidate_centerlines = [np.full((seq_len, 2), None)]
candidate_nt_distances = [np.full((seq_len, 2), None)]
# Get NT distance for oracle centerline
oracle_nt_dist = get_nt_distance(agent_xy,
oracle_centerline,
viz=False)
map_feature_helpers = {
"ORACLE_CENTERLINE": oracle_centerline,
"CANDIDATE_CENTERLINES": candidate_centerlines,
"CANDIDATE_NT_DISTANCES": candidate_nt_distances,
}
return oracle_nt_dist, map_feature_helpers
def compute_map_features(
self, agent_track: np.ndarray, obs_len: int, seq_len: int,
raw_data_format: Dict[str, int], mode: str,
avm: ArgoverseMap) -> Tuple[np.ndarray, Dict[str, Any]]:
"""Compute map based features for the given sequence.
If the mode is test, oracle_nt_dist will be empty, candidate_nt_dist will be populated.
If the mode is train/val, oracle_nt_dist will be populated, candidate_nt_dist will be empty.
Args:
agent_track : Data for the agent track
obs_len : Length of observed trajectory
seq_len : Length of the sequence
raw_data_format : Format of the sequence
mode: train/val/test mode
Returns:
oracle_nt_dist (numpy array): normal and tangential distances for oracle centerline
map_feature_helpers (dict): Dictionary containing helpers for map features
"""
obs_pred_lanes = []
unique_segments_future = []
unique_segments_past = []
curr_lane_candidates = []
# Get observed 2 secs of the agent
agent_xy = agent_track[:, [raw_data_format["X"], raw_data_format["Y"]
]].astype("float")
agent_track_obs = agent_track[:obs_len]
agent_xy_obs = agent_track_obs[:, [
raw_data_format["X"], raw_data_format["Y"]
]].astype("float")
# Get API for Argo Dataset map
city_name = agent_track[0, raw_data_format["CITY_NAME"]]
# Get candidate centerlines using observed trajectory
if mode == "test":
oracle_centerline = np.full((seq_len, 2), None)
oracle_nt_dist = np.full((seq_len, 2), None)
candidate_centerlines = self.get_candidate_centerlines_for_trajectory(
agent_xy_obs,
city_name,
avm,
viz=False,
max_search_radius=self._MAX_SEARCH_RADIUS_CENTERLINES,
seq_len=seq_len,
max_candidates=self._MAX_CENTERLINE_CANDIDATES_TEST,
)
# Get nt distance for the entire trajectory using candidate centerlines
candidate_nt_distances = []
for candidate_centerline in candidate_centerlines:
candidate_nt_distance = np.full((seq_len, 2), None)
candidate_nt_distance[:obs_len] = get_nt_distance(
agent_xy_obs, candidate_centerline)
candidate_nt_distances.append(candidate_nt_distance)
elif mode == "compute_all":
# Get oracle centerline
oracle_centerline = self.get_candidate_centerlines_for_trajectory(
agent_xy,
city_name,
avm,
viz=False,
max_search_radius=self._MAX_SEARCH_RADIUS_CENTERLINES,
seq_len=seq_len,
mode="train",
)[0]
# Get NT distance for oracle centerline
oracle_nt_dist = get_nt_distance(agent_xy,
oracle_centerline,
viz=False)
# Get candidate centerl = []ines
candidate_centerlines, obs_pred_lanes, unique_segments_future, unique_segments_past, curr_lane_candidates = self.get_candidate_centerlines_for_trajectory(
agent_xy_obs,
city_name,
avm,
viz=False,
max_search_radius=self._MAX_SEARCH_RADIUS_CENTERLINES,
seq_len=seq_len,
max_candidates=self._MAX_CENTERLINE_CANDIDATES_TEST,
mode=mode)
# Get nt distance for the entire trajectory using candidate centerlines
candidate_nt_distances = []
for candidate_centerline in candidate_centerlines:
candidate_nt_distance = np.full((seq_len, 2), None)
candidate_nt_distance[:obs_len] = get_nt_distance(
agent_xy_obs, candidate_centerline)
candidate_nt_distances.append(candidate_nt_distance)
elif mode == "lanes_only":
# Get oracle centerline
oracle_centerline = self.get_candidate_centerlines_for_trajectory(
agent_xy,
city_name,
avm,
viz=False,
max_search_radius=self._MAX_SEARCH_RADIUS_CENTERLINES,
seq_len=seq_len,
mode="train",
)[0]
# Not computing oracle nt_distances
oracle_nt_dist = np.full((seq_len, 2), None)
# Get candidate centerl = []ines
candidate_centerlines, obs_pred_lanes, unique_segments_future, unique_segments_past, curr_lane_candidates = self.get_candidate_centerlines_for_trajectory(
agent_xy_obs,
city_name,
avm,
viz=False,
max_search_radius=self._MAX_SEARCH_RADIUS_CENTERLINES,
seq_len=seq_len,
max_candidates=self._MAX_CENTERLINE_CANDIDATES_TEST,
mode=mode)
# Not computing candidate nt_distances
candidate_nt_distances = []
else:
oracle_centerline = self.get_candidate_centerlines_for_trajectory(
agent_xy,
city_name,
avm,
viz=False,
max_search_radius=self._MAX_SEARCH_RADIUS_CENTERLINES,
seq_len=seq_len,
mode=mode,
)[0]
candidate_centerlines = [np.full((seq_len, 2), None)]
candidate_nt_distances = [np.full((seq_len, 2), None)]
# Get NT distance for oracle centerline
oracle_nt_dist = get_nt_distance(agent_xy,
oracle_centerline,
viz=False)
# ADD A LOOP HERE TO GO OVER THE CANDIDATE CENTERLINES. CURRENTLY WE DO IT ONLY FOR ONE LINE.
# here, do the lead and following agent assignments
initial_agent_ids = df[
"TRACK_ID"].values # All the track_ids (unique actor identifier in .csv dataset)
initial_agent_ids = list(
dict.fromkeys(initial_agent_ids)) # Remove duplicates from IDs
removed_agent_ids = initial_agent_ids
# Remove the current agent we are calculating the features for
removed_agent_ids.remove(track_id)
# here need to loop thru candidate centerlines to return for every candidate centerline.
# need to add normal dist. threshold bc the lead vehicle might be returned even tho it is far from the lane
# since algorthm returns the closest currently
# need to add threshold for distance to the agent.
# need to convert to desired standard format that is (centerline, feature)
leading_actor = []
following_actor = []
leading_actor_dist = []
following_actor_dist = []
#print(candidate_centerlines)
c_line_idx = 0
# for c_line in candidate_centerlines:
# leading_actor_idx, following_actor_idx, \
# leading_actor_dist_idx, following_actor_dist_idx = \
# self.get_lead_and_follow(df, obs_len, agent_track, removed_agent_ids, \
# c_line, raw_data_format)
# leading_actor.append(leading_actor_idx)
# following_actor.append(following_actor_idx)
# leading_actor_dist.append(leading_actor_dist_idx)
# following_actor_dist.append(following_actor_dist_idx)
# c_line_idx += 1
map_feature_helpers = {
"LEADING_VEHICLES": leading_actor,
"FOLLOWING_VEHICLES": following_actor,
"LEADING_DISTANCES": leading_actor_dist,
"FOLLOWING_DISTANCES": following_actor_dist,
"ORACLE_CENTERLINE": oracle_centerline,
"CANDIDATE_CENTERLINES": candidate_centerlines,
"CANDIDATE_NT_DISTANCES": candidate_nt_distances,
"CANDIDATE_LANE_SEGMENTS": obs_pred_lanes,
"LANE_SEGMENTS_IN_BUBBLE": curr_lane_candidates,
"LANE_SEGMENTS_IN_FRONT": unique_segments_future,
"LANE_SEGMENTS_IN_BACK": unique_segments_past
}
return oracle_nt_dist, map_feature_helpers