From 69a3d7a4e48842db0021f06fd25226862fdb085b Mon Sep 17 00:00:00 2001 From: Quanyi Li Date: Sat, 24 Feb 2024 13:34:31 +0000 Subject: [PATCH] Fix nuscenes (#67) * fix bug * fix a potential bug --- scenarionet/converter/nuscenes/utils.py | 70 +++++++++---------------- 1 file changed, 24 insertions(+), 46 deletions(-) diff --git a/scenarionet/converter/nuscenes/utils.py b/scenarionet/converter/nuscenes/utils.py index 5873f9d..7f9eb95 100644 --- a/scenarionet/converter/nuscenes/utils.py +++ b/scenarionet/converter/nuscenes/utils.py @@ -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