Fix nuscenes (#67)

* fix bug

* fix a potential bug
This commit is contained in:
Quanyi Li
2024-02-24 13:34:31 +00:00
committed by GitHub
parent 6cda061ed8
commit 69a3d7a4e4

View File

@@ -256,17 +256,7 @@ def get_tracks_from_frames(nuscenes: NuScenes, scene_info, frames, num_to_interp
interpolate_tracks[id]["state"][k] = interpolate(v, track["state"]["valid"], new_valid)
# if id == "ego":
# ego is valid all time, so we can calculate the velocity in this way
# Normalize place all object to (0,0)
map_center = np.array(interpolate_tracks["ego"]["state"]["position"][0])
map_center[-1] = 0
normalized_ret = {}
for id, track, in interpolate_tracks.items():
pos = track["state"]["position"] - map_center
track["state"]["position"] = np.asarray(pos)
normalized_ret[id] = track
return normalized_ret, map_center
return interpolate_tracks
def get_map_features(scene_info, nuscenes: NuScenes, map_center, radius=500, points_distance=1, only_lane=False):
@@ -323,10 +313,7 @@ def get_map_features(scene_info, nuscenes: NuScenes, map_center, radius=500, poi
for idx, boundary in enumerate(boundaries[0]):
block_points = np.array(list(i for i in zip(boundary.coords.xy[0], boundary.coords.xy[1])))
id = "boundary_{}".format(idx)
ret[id] = {
SD.TYPE: MetaDriveType.LINE_SOLID_SINGLE_WHITE,
SD.POLYLINE: block_points - np.asarray(map_center)[:2]
}
ret[id] = {SD.TYPE: MetaDriveType.LINE_SOLID_SINGLE_WHITE, SD.POLYLINE: block_points}
# broken line
for id in map_objs["lane_divider"]:
@@ -334,7 +321,7 @@ def get_map_features(scene_info, nuscenes: NuScenes, map_center, radius=500, poi
assert line_info["token"] == id
line = map_api.extract_line(line_info["line_token"]).coords.xy
line = np.asarray([[line[0][i], line[1][i]] for i in range(len(line[0]))])
ret[id] = {SD.TYPE: MetaDriveType.LINE_BROKEN_SINGLE_WHITE, SD.POLYLINE: line - np.asarray(map_center)[:2]}
ret[id] = {SD.TYPE: MetaDriveType.LINE_BROKEN_SINGLE_WHITE, SD.POLYLINE: line}
# solid line
for id in map_objs["road_divider"]:
@@ -342,7 +329,7 @@ def get_map_features(scene_info, nuscenes: NuScenes, map_center, radius=500, poi
assert line_info["token"] == id
line = map_api.extract_line(line_info["line_token"]).coords.xy
line = np.asarray([[line[0][i], line[1][i]] for i in range(len(line[0]))])
ret[id] = {SD.TYPE: MetaDriveType.LINE_SOLID_SINGLE_YELLOW, SD.POLYLINE: line - np.asarray(map_center)[:2]}
ret[id] = {SD.TYPE: MetaDriveType.LINE_SOLID_SINGLE_YELLOW, SD.POLYLINE: line}
# crosswalk
for id in map_objs["ped_crossing"]:
@@ -352,7 +339,7 @@ def get_map_features(scene_info, nuscenes: NuScenes, map_center, radius=500, poi
boundary_polygon = np.asarray([[boundary[0][i], boundary[1][i]] for i in range(len(boundary[0]))])
ret[id] = {
SD.TYPE: MetaDriveType.CROSSWALK,
SD.POLYGON: boundary_polygon - np.asarray(map_center)[:2],
SD.POLYGON: boundary_polygon,
}
# walkway
@@ -363,7 +350,7 @@ def get_map_features(scene_info, nuscenes: NuScenes, map_center, radius=500, poi
boundary_polygon = np.asarray([[boundary[0][i], boundary[1][i]] for i in range(len(boundary[0]))])
ret[id] = {
SD.TYPE: MetaDriveType.BOUNDARY_SIDEWALK,
SD.POLYGON: boundary_polygon - np.asarray(map_center)[:2],
SD.POLYGON: boundary_polygon,
}
# normal lane
@@ -375,9 +362,8 @@ def get_map_features(scene_info, nuscenes: NuScenes, map_center, radius=500, poi
# boundary_polygon += [[boundary[0][i], boundary[1][i]] for i in range(len(boundary[0]))]
ret[id] = {
SD.TYPE: MetaDriveType.LANE_SURFACE_STREET,
SD.POLYLINE: np.asarray(discretize_lane(map_api.arcline_path_3[id], resolution_meters=points_distance)) -
np.asarray(map_center),
SD.POLYGON: boundary_polygon - np.asarray(map_center)[:2],
SD.POLYLINE: np.asarray(discretize_lane(map_api.arcline_path_3[id], resolution_meters=points_distance)),
SD.POLYGON: boundary_polygon,
SD.ENTRY: map_api.get_incoming_lane_ids(id),
SD.EXIT: map_api.get_outgoing_lane_ids(id),
SD.LEFT_NEIGHBORS: [],
@@ -393,8 +379,7 @@ def get_map_features(scene_info, nuscenes: NuScenes, map_center, radius=500, poi
# boundary_polygon += [[boundary[0][i], boundary[1][i], 0.] for i in range(len(boundary[0]))]
ret[id] = {
SD.TYPE: MetaDriveType.LANE_SURFACE_UNSTRUCTURE,
SD.POLYLINE: np.asarray(discretize_lane(map_api.arcline_path_3[id], resolution_meters=points_distance)) -
np.asarray(map_center),
SD.POLYLINE: np.asarray(discretize_lane(map_api.arcline_path_3[id], resolution_meters=points_distance)),
# SD.POLYGON: boundary_polygon,
"speed_limit_kmh": 100,
SD.ENTRY: map_api.get_incoming_lane_ids(id),
@@ -409,7 +394,7 @@ def get_map_features(scene_info, nuscenes: NuScenes, map_center, radius=500, poi
# boundary_polygon = np.asarray([[boundary[0][i], boundary[1][i]] for i in range(len(boundary[0]))])
# ret[id] = {
# SD.TYPE: MetaDriveType.STOP_LINE,
# SD.POLYGON: boundary_polygon - np.asarray(map_center)[:2],
# SD.POLYGON: boundary_polygon ,
# }
# 'stop_line',
@@ -469,32 +454,25 @@ def convert_nuscenes_scenario(
result[SD.METADATA]["sample_rate"] = scenario_log_interval
result[SD.METADATA][SD.TIMESTEP] = np.arange(0., (len(frames) - 1) * 0.5 + 0.1, 0.1)
# interpolating to 0.1s interval
result[SD.TRACKS], map_center = get_tracks_from_frames(nuscenes, scene_info, frames, num_to_interpolate=5)
result[SD.TRACKS] = get_tracks_from_frames(nuscenes, scene_info, frames, num_to_interpolate=5)
result[SD.METADATA][SD.SDC_ID] = "ego"
# No traffic light in nuscenes at this stage
result[SD.DYNAMIC_MAP_STATES] = {}
# track_to_predict = result[SD.TRACKS][instance_token]
# result[SD.METADATA]["tracks_to_predict"] = {
# instance_token: {
# "track_index": list(result[SD.TRACKS].keys()).index(instance_token),
# "track_id": instance_token,
# "difficulty": 0,
# "object_type": track_to_predict['type']
# }
# }
# map
result[SD.MAP_FEATURES] = get_map_features(scene_info, nuscenes, map_center, map_radius, only_lane=only_lane)
# add back map center
map_center = map_center[np.newaxis]
for k, v in result[SD.TRACKS].items():
v['state']['position'] += map_center
if prediction:
track_to_predict = result[SD.TRACKS][instance_token]
result[SD.METADATA]["tracks_to_predict"] = {
instance_token: {
"track_index": list(result[SD.TRACKS].keys()).index(instance_token),
"track_id": instance_token,
"difficulty": 0,
"object_type": track_to_predict['type']
}
}
for k, v in result[SD.MAP_FEATURES].items():
if 'polygon' in v:
v['polygon'] += map_center[:, :v['polygon'].shape[-1]]
else:
v['polyline'] += map_center[:, :v['polyline'].shape[-1]]
# map
map_center = np.array(result[SD.TRACKS]["ego"]["state"]["position"][0])
result[SD.MAP_FEATURES] = get_map_features(scene_info, nuscenes, map_center, map_radius, only_lane=only_lane)
del frames_scene_info
del frames
del scene_info