diff --git a/scenarionet/converter/argoverse2/utils.py b/scenarionet/converter/argoverse2/utils.py index 0e6dda7..5131823 100644 --- a/scenarionet/converter/argoverse2/utils.py +++ b/scenarionet/converter/argoverse2/utils.py @@ -3,7 +3,8 @@ import logging import tqdm from scenarionet.converter.utils import mph_to_kmh - +import geopandas as gpd +from shapely.ops import unary_union logger = logging.getLogger(__name__) import numpy as np @@ -16,6 +17,7 @@ from scenarionet.converter.argoverse2.type import get_traffic_obj_type, get_lane from av2.datasets.motion_forecasting import scenario_serialization from av2.map.map_api import ArgoverseStaticMap from typing import Final +from shapely.geometry import Point, Polygon _ESTIMATED_VEHICLE_LENGTH_M: Final[float] = 4.0 _ESTIMATED_VEHICLE_WIDTH_M: Final[float] = 2.0 @@ -151,11 +153,19 @@ def extract_map_features(map_features): ret[lane_id] = center - # for edge in vector_drivable_areas: - # bound = dict() - # bound["type"] = MetaDriveType.BOUNDARY_LINE - # bound["polyline"] = edge.xyz.astype(np.float32) - # ret[str(edge.id)] = bound + polygons = [] + for polygon in vector_drivable_areas: + # convert to shapely polygon + points = polygon.area_boundary + polygons.append(Polygon([(p.x, p.y) for p in points])) + + polygons = [geom if geom.is_valid else geom.buffer(0) for geom in polygons] + boundaries = gpd.GeoSeries(unary_union(polygons)).boundary.explode(index_parts=True) + 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]))) + for i in range(0, len(block_points), 20): + id = f'boundary_{idx}{i}' + ret[id] = {SD.TYPE: MetaDriveType.LINE_SOLID_SINGLE_WHITE, SD.POLYLINE: block_points[i:i + 20]} for cross in ped_crossings: bound = dict() diff --git a/scenarionet/converter/nuscenes/utils.py b/scenarionet/converter/nuscenes/utils.py index 710a9ab..84392f2 100644 --- a/scenarionet/converter/nuscenes/utils.py +++ b/scenarionet/converter/nuscenes/utils.py @@ -474,7 +474,15 @@ def convert_nuscenes_scenario( # 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) del frames_scene_info