-
Notifications
You must be signed in to change notification settings - Fork 20
/
data_process.py
192 lines (150 loc) · 8.85 KB
/
data_process.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
import os
import argparse
from tqdm import tqdm
from common_utils import *
from GameFormer.data_utils import *
import matplotlib.pyplot as plt
from nuplan.planning.utils.multithreading.worker_parallel import SingleMachineParallelExecutor
from nuplan.planning.scenario_builder.scenario_filter import ScenarioFilter
from nuplan.planning.scenario_builder.nuplan_db.nuplan_scenario_builder import NuPlanScenarioBuilder
from nuplan.planning.scenario_builder.nuplan_db.nuplan_scenario_utils import ScenarioMapping
# define data processor
class DataProcessor(object):
def __init__(self, scenarios):
self._scenarios = scenarios
self.past_time_horizon = 2 # [seconds]
self.num_past_poses = 10 * self.past_time_horizon
self.future_time_horizon = 8 # [seconds]
self.num_future_poses = 10 * self.future_time_horizon
self.num_agents = 20
self._map_features = ['LANE', 'ROUTE_LANES', 'CROSSWALK'] # name of map features to be extracted.
self._max_elements = {'LANE': 40, 'ROUTE_LANES': 10, 'CROSSWALK': 5} # maximum number of elements to extract per feature layer.
self._max_points = {'LANE': 50, 'ROUTE_LANES': 50, 'CROSSWALK': 30} # maximum number of points per feature to extract per feature layer.
self._radius = 60 # [m] query radius scope relative to the current pose.
self._interpolation_method = 'linear' # Interpolation method to apply when interpolating to maintain fixed size map elements.
def get_ego_agent(self):
self.anchor_ego_state = self.scenario.initial_ego_state
past_ego_states = self.scenario.get_ego_past_trajectory(
iteration=0, num_samples=self.num_past_poses, time_horizon=self.past_time_horizon
)
sampled_past_ego_states = list(past_ego_states) + [self.anchor_ego_state]
past_ego_states_tensor = sampled_past_ego_states_to_tensor(sampled_past_ego_states)
past_time_stamps = list(
self.scenario.get_past_timestamps(
iteration=0, num_samples=self.num_past_poses, time_horizon=self.past_time_horizon
)
) + [self.scenario.start_time]
past_time_stamps_tensor = sampled_past_timestamps_to_tensor(past_time_stamps)
return past_ego_states_tensor, past_time_stamps_tensor
def get_neighbor_agents(self):
present_tracked_objects = self.scenario.initial_tracked_objects.tracked_objects
past_tracked_objects = [
tracked_objects.tracked_objects
for tracked_objects in self.scenario.get_past_tracked_objects(
iteration=0, time_horizon=self.past_time_horizon, num_samples=self.num_past_poses
)
]
sampled_past_observations = past_tracked_objects + [present_tracked_objects]
past_tracked_objects_tensor_list, past_tracked_objects_types = \
sampled_tracked_objects_to_tensor_list(sampled_past_observations)
return past_tracked_objects_tensor_list, past_tracked_objects_types
def get_map(self):
ego_state = self.scenario.initial_ego_state
ego_coords = Point2D(ego_state.rear_axle.x, ego_state.rear_axle.y)
route_roadblock_ids = self.scenario.get_route_roadblock_ids()
traffic_light_data = self.scenario.get_traffic_light_status_at_iteration(0)
coords, traffic_light_data = get_neighbor_vector_set_map(
self.map_api, self._map_features, ego_coords, self._radius, route_roadblock_ids, traffic_light_data
)
vector_map = map_process(ego_state.rear_axle, coords, traffic_light_data, self._map_features,
self._max_elements, self._max_points, self._interpolation_method)
return vector_map
def get_ego_agent_future(self):
current_absolute_state = self.scenario.initial_ego_state
trajectory_absolute_states = self.scenario.get_ego_future_trajectory(
iteration=0, num_samples=self.num_future_poses, time_horizon=self.future_time_horizon
)
# Get all future poses of the ego relative to the ego coordinate system
trajectory_relative_poses = convert_absolute_to_relative_poses(
current_absolute_state.rear_axle, [state.rear_axle for state in trajectory_absolute_states]
)
return trajectory_relative_poses
def get_neighbor_agents_future(self, agent_index):
current_ego_state = self.scenario.initial_ego_state
present_tracked_objects = self.scenario.initial_tracked_objects.tracked_objects
# Get all future poses of of other agents
future_tracked_objects = [
tracked_objects.tracked_objects
for tracked_objects in self.scenario.get_future_tracked_objects(
iteration=0, time_horizon=self.future_time_horizon, num_samples=self.num_future_poses
)
]
sampled_future_observations = [present_tracked_objects] + future_tracked_objects
future_tracked_objects_tensor_list, _ = sampled_tracked_objects_to_tensor_list(sampled_future_observations)
agent_futures = agent_future_process(current_ego_state, future_tracked_objects_tensor_list, self.num_agents, agent_index)
return agent_futures
def plot_scenario(self, data):
# Create map layers
create_map_raster(data['lanes'], data['crosswalks'], data['route_lanes'])
# Create agent layers
create_ego_raster(data['ego_agent_past'][-1])
create_agents_raster(data['neighbor_agents_past'][:, -1])
# Draw past and future trajectories
draw_trajectory(data['ego_agent_past'], data['neighbor_agents_past'])
draw_trajectory(data['ego_agent_future'], data['neighbor_agents_future'])
plt.gca().set_aspect('equal')
plt.tight_layout()
plt.show()
def save_to_disk(self, dir, data):
np.savez(f"{dir}/{data['map_name']}_{data['token']}.npz", **data)
def work(self, save_dir, debug=False):
for scenario in tqdm(self._scenarios):
map_name = scenario._map_name
token = scenario.token
self.scenario = scenario
self.map_api = scenario.map_api
# get agent past tracks
ego_agent_past, time_stamps_past = self.get_ego_agent()
neighbor_agents_past, neighbor_agents_types = self.get_neighbor_agents()
ego_agent_past, neighbor_agents_past, neighbor_indices = \
agent_past_process(ego_agent_past, time_stamps_past, neighbor_agents_past, neighbor_agents_types, self.num_agents)
# get vector set map
vector_map = self.get_map()
# get agent future tracks
ego_agent_future = self.get_ego_agent_future()
neighbor_agents_future = self.get_neighbor_agents_future(neighbor_indices)
# gather data
data = {"map_name": map_name, "token": token, "ego_agent_past": ego_agent_past, "ego_agent_future": ego_agent_future,
"neighbor_agents_past": neighbor_agents_past, "neighbor_agents_future": neighbor_agents_future}
data.update(vector_map)
# visualization
if debug:
self.plot_scenario(data)
# save to disk
self.save_to_disk(save_dir, data)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Data Processing')
parser.add_argument('--data_path', type=str, help='path to raw data')
parser.add_argument('--map_path', type=str, help='path to map data')
parser.add_argument('--save_path', type=str, help='path to save processed data')
parser.add_argument('--scenarios_per_type', type=int, default=1000, help='number of scenarios per type')
parser.add_argument('--total_scenarios', default=None, help='limit total number of scenarios')
parser.add_argument('--shuffle_scenarios', type=bool, default=False, help='shuffle scenarios')
parser.add_argument('--debug', action="store_true", help='if visualize the data output', default=False)
args = parser.parse_args()
# create save folder
os.makedirs(args.save_path, exist_ok=True)
# get scenarios
map_version = "nuplan-maps-v1.0"
sensor_root = None
db_files = None
scenario_mapping = ScenarioMapping(scenario_map=get_scenario_map(), subsample_ratio_override=0.5)
builder = NuPlanScenarioBuilder(args.data_path, args.map_path, sensor_root, db_files, map_version, scenario_mapping=scenario_mapping)
scenario_filter = ScenarioFilter(*get_filter_parameters(args.scenarios_per_type, args.total_scenarios, args.shuffle_scenarios))
worker = SingleMachineParallelExecutor(use_process_pool=True)
scenarios = builder.get_scenarios(scenario_filter, worker)
print(f"Total number of scenarios: {len(scenarios)}")
# process data
del worker, builder, scenario_filter, scenario_mapping
processor = DataProcessor(scenarios)
processor.work(args.save_path, debug=args.debug)