diff --git a/.github/workflows/ci-base-tests-linux.yml b/.github/workflows/ci-base-tests-linux.yml index e3a02a628f..e68ae05ddc 100644 --- a/.github/workflows/ci-base-tests-linux.yml +++ b/.github/workflows/ci-base-tests-linux.yml @@ -29,7 +29,7 @@ jobs: . ${{env.venv_dir}}/bin/activate pip install --upgrade pip pip install --upgrade wheel - pip install -e .[camera_obs,opendrive,rllib,test,test_notebook,torch,train,gym] + pip install -e .[camera_obs,opendrive,rllib,test,test_notebook,torch,train,gym,argoverse] - name: Run smoke tests run: | . ${{env.venv_dir}}/bin/activate @@ -41,6 +41,7 @@ jobs: -n auto \ --ignore-glob="**/ros.py" \ --ignore-glob="**/waymo_map.py" \ + --ignore-glob="**/argoverse_map.py" \ ${{matrix.tests}} \ --ignore=./smarts/core/tests/test_smarts_memory_growth.py \ --ignore=./smarts/core/tests/test_env_frame_rate.py \ diff --git a/.github/workflows/ci-base-tests-mac.yml b/.github/workflows/ci-base-tests-mac.yml index f4deffffcc..d10098989b 100644 --- a/.github/workflows/ci-base-tests-mac.yml +++ b/.github/workflows/ci-base-tests-mac.yml @@ -47,7 +47,7 @@ jobs: pip install --upgrade pip pip install --upgrade wheel pip install -r utils/setup/mac_requirements.txt - pip install -e .[camera_obs,opendrive,rllib,test,test_notebook,torch,train] + pip install -e .[camera_obs,opendrive,rllib,test,test_notebook,torch,train,argoverse] - name: Run smoke tests run: | . ${{env.venv_dir}}/bin/activate @@ -58,6 +58,7 @@ jobs: --doctest-modules \ -n auto \ --ignore-glob="**/waymo_map.py" \ + --ignore-glob="**/argoverse_map.py" \ ${{matrix.tests}} \ --ignore=./smarts/core/tests/test_smarts_memory_growth.py \ --ignore=./smarts/env/tests/test_benchmark.py \ diff --git a/CHANGELOG.md b/CHANGELOG.md index 7dc06725ce..983fa89900 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ Copy and pasting the git commit messages is __NOT__ enough. ## [Unreleased] ### Added +- Added support for the [Argoverse 2 Motion Forecasting Dataset](https://www.argoverse.org/av2.html#forecasting-link) (see `scenarios/argoverse`) ### Changed ### Deprecated ### Fixed diff --git a/Makefile b/Makefile index d05f63c478..f424211507 100644 --- a/Makefile +++ b/Makefile @@ -10,6 +10,9 @@ test: build-all-scenarios -n `expr \( \`nproc\` \/ 2 \& \`nproc\` \> 3 \) \| 2` \ --nb-exec-timeout 65536 \ ./examples/tests ./smarts/env ./envision ./smarts/core ./smarts/sstudio ./tests \ + --ignore=./smarts/core/waymo_map.py \ + --ignore=./smarts/core/argoverse_map.py \ + --ignore=./smarts/core/tests/test_argoverse.py \ --ignore=./smarts/core/tests/test_smarts_memory_growth.py \ --ignore=./smarts/core/tests/test_env_frame_rate.py \ --ignore=./smarts/core/tests/test_notebook.py \ diff --git a/docs/_static/argoverse-replay.gif b/docs/_static/argoverse-replay.gif new file mode 100644 index 0000000000..ab08f55ae9 Binary files /dev/null and b/docs/_static/argoverse-replay.gif differ diff --git a/docs/ecosystem/argoverse.rst b/docs/ecosystem/argoverse.rst new file mode 100644 index 0000000000..ba776c21bd --- /dev/null +++ b/docs/ecosystem/argoverse.rst @@ -0,0 +1,32 @@ +.. _argoverse: + +Argoverse 2 +=========== + +SMARTS supports loading and replaying scenarios from the **Argoverse 2 Motion Forecasting Dataset**. See the Argoverse `website `_ for more info and download instructions. + +An example SMARTS scenario is located `here `_. After downloading the dataset, modify the ``scenario_id`` and ``scenario_path`` variables to point to the desired Argoverse scenario: + +.. code-block:: python + + # scenario_path is a directory with the following structure: + # /path/to/dataset/{scenario_id} + # ├── log_map_archive_{scenario_id}.json + # └── scenario_{scenario_id}.parquet + + scenario_id = "0000b6ab-e100-4f6b-aee8-b520b57c0530" + scenario_path = Path("/home/user/argoverse/train/") / scenario_id + +Make sure you've installed the required dependencies: + +.. code-block:: sh + + $ pip install -e .[argoverse] + +You can then run any of the examples with this scenario: + +.. code-block:: sh + + $ scl run --envision examples/egoless.py scenarios/argoverse + +.. image:: /_static/argoverse-replay.gif diff --git a/docs/index.rst b/docs/index.rst index 5563aab550..0bbbe10da2 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -79,6 +79,7 @@ If you use SMARTS in your research, please cite the `paper =1.2.1 +argoverse = + av2>=0.2.1 [aliases] test=pytest diff --git a/smarts/core/argoverse_map.py b/smarts/core/argoverse_map.py new file mode 100644 index 0000000000..0e570ec884 --- /dev/null +++ b/smarts/core/argoverse_map.py @@ -0,0 +1,1127 @@ +# Copyright (C) 2022. Huawei Technologies Co., Ltd. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import heapq +import logging +import random +import time +from functools import lru_cache +from pathlib import Path +from typing import Dict, List, Optional, Sequence, Set, Tuple + +import numpy as np +import rtree +from cached_property import cached_property +from shapely.geometry import Point as SPoint +from shapely.geometry import Polygon + +from smarts.core.coordinates import BoundingBox, Heading, Point, Pose, RefLinePoint +from smarts.core.lanepoints import LanePoints, LinkedLanePoint +from smarts.core.road_map import RoadMap, RoadMapWithCaches, Waypoint +from smarts.core.route_cache import RouteWithCache +from smarts.core.utils.glb import make_map_glb, make_road_line_glb +from smarts.core.utils.math import ( + inplace_unwrap, + line_intersect_vectorized, + radians_to_vec, + vec_2d, +) +from smarts.sstudio.types import MapSpec + +try: + from av2.geometry.interpolate import interp_arc + from av2.map.lane_segment import LaneMarkType, LaneSegment + from av2.map.map_api import ArgoverseStaticMap +except: + raise ImportError( + "You may not have installed the [argoverse] dependencies required for using Argoverse 2 maps with SMARTS. Install it first using the command `pip install -e .[argoverse]` at the source directory." + "" + ) + + +class ArgoverseMap(RoadMapWithCaches): + """A road map for an Argoverse 2 scenario.""" + + DEFAULT_LANE_SPEED = 16.67 # m/s + + LANE_MARKINGS = frozenset( + { + LaneMarkType.DASH_SOLID_WHITE, + LaneMarkType.DASHED_WHITE, + LaneMarkType.DOUBLE_SOLID_WHITE, + LaneMarkType.DOUBLE_DASH_WHITE, + LaneMarkType.SOLID_WHITE, + LaneMarkType.SOLID_DASH_WHITE, + } + ) + + ROAD_MARKINGS = frozenset( + { + LaneMarkType.DASH_SOLID_YELLOW, + LaneMarkType.DASHED_YELLOW, + LaneMarkType.DOUBLE_SOLID_YELLOW, + LaneMarkType.DOUBLE_DASH_YELLOW, + LaneMarkType.SOLID_YELLOW, + LaneMarkType.SOLID_DASH_YELLOW, + LaneMarkType.SOLID_BLUE, + } + ) + + def __init__(self, map_spec: MapSpec, avm: ArgoverseStaticMap): + super().__init__() + self._log = logging.getLogger(self.__class__.__name__) + self._avm = avm + self._argoverse_scenario_id = avm.log_id + self._map_spec = map_spec + self._surfaces = dict() + self._lanes: Dict[str, ArgoverseMap.Lane] = dict() + self._roads: Dict[str, ArgoverseMap.Road] = dict() + self._features = dict() + self._lane_rtree = None + self._load_map_data() + self._waypoints_cache = ArgoverseMap._WaypointsCache() + if map_spec.lanepoint_spacing is not None: + assert map_spec.lanepoint_spacing > 0 + self._lanepoints = LanePoints.from_argoverse( + self, spacing=map_spec.lanepoint_spacing + ) + + @classmethod + def from_spec(cls, map_spec: MapSpec): + """Generate a road map from the given specification.""" + scenario_dir = Path(map_spec.source) + scenario_id = scenario_dir.stem + map_path = scenario_dir / f"log_map_archive_{scenario_id}.json" + + if not map_path.exists(): + logging.warning(f"Map not found: {map_path}") + return None + + avm = ArgoverseStaticMap.from_json(map_path) + assert avm.log_id == scenario_id, "Loaded map ID does not match expected ID" + return cls(map_spec, avm) + + def _compute_lane_intersections(self): + intersections: Dict[str, Set[str]] = dict() + + lane_ids_todo = [lane_id for lane_id in self._lanes.keys()] + + # Build rtree + lane_rtree = rtree.index.Index() + lane_rtree.interleaved = True + bboxes = dict() + for idx, lane_id in enumerate(lane_ids_todo): + # Using the centerline here is much faster than using the lane polygon + lane_pts = self._lanes[lane_id]._centerline + bbox = ( + np.amin(lane_pts[:, 0]), + np.amin(lane_pts[:, 1]), + np.amax(lane_pts[:, 0]), + np.amax(lane_pts[:, 1]), + ) + bboxes[lane_id] = bbox + lane_rtree.add(idx, bbox) + + for lane_id in lane_ids_todo: + lane = self._lanes[lane_id] + lane_intersections = intersections.setdefault(lane_id, set()) + + # Filter out any lanes that don't intersect this lane's bbox + indicies = lane_rtree.intersection(bboxes[lane_id]) + + # Filter out any other lanes we don't want to check against + lanes_to_test = [] + for idx in indicies: + cand_id = lane_ids_todo[idx] + if cand_id == lane_id: + continue + # Skip intersections we've already computed + if cand_id in lane_intersections: + continue + # ... and sub-lanes of the same original lane + cand_lane = self._lanes[cand_id] + # Don't check intersection with incoming/outgoing lanes + if cand_lane in lane.incoming_lanes or cand_lane in lane.outgoing_lanes: + continue + # ... or lanes in same road (TAI?) + if lane.road == cand_lane.road: + continue + lanes_to_test.append(cand_id) + if not lanes_to_test: + continue + + # Main loop -- check each segment of the lane polyline against the + # polyline of each candidate lane (--> algorithm is O(l^2) + line1 = lane._centerline + for cand_id in lanes_to_test: + line2 = np.array(self._lanes[cand_id]._centerline) + C = np.roll(line2, 0, axis=0)[:-1] + D = np.roll(line2, -1, axis=0)[:-1] + for i in range(len(line1) - 1): + a = line1[i] + b = line1[i + 1] + if line_intersect_vectorized(a, b, C, D): + lane_intersections.add(cand_id) + intersections.setdefault(cand_id, set()).add(lane_id) + break + + for lane_id, intersect_ids in intersections.items(): + self._lanes[lane_id]._intersections = [ + self.lane_by_id(id) for id in intersect_ids + ] + + def _load_map_data(self): + start = time.time() + + all_ids = set(self._avm.get_scenario_lane_segment_ids()) + processed_ids = set() + for lane_seg in self._avm.get_scenario_lane_segments(): + # If this is a rightmost lane, create a road with its neighbours + if lane_seg.right_neighbor_id is None: + neighbours: List[int] = [] + cur_seg = lane_seg + while True: + left_mark = cur_seg.left_lane_marking.mark_type + left_id = cur_seg.left_neighbor_id + if ( + left_id is not None + and left_mark == LaneMarkType.DASHED_WHITE + and left_id in self._avm.vector_lane_segments + ): + # There is a valid lane to the left, so add it and continue + left_seg = self._avm.vector_lane_segments[left_id] + + # Edge case: sometimes there can be a cycle (2 lanes can have each other as their left neighbour) + if left_seg.left_neighbor_id == cur_seg.id: + break + + cur_seg = left_seg + neighbours.append(left_id) + else: + break # This is the leftmost lane in the road, so stop + + # Create the lane objects + road_id = "road" + lanes = [] + for index, seg_id in enumerate([lane_seg.id] + neighbours): + road_id += f"-{seg_id}" + lane_id = f"lane-{seg_id}" + seg = self._avm.vector_lane_segments[seg_id] + lane = ArgoverseMap.Lane(self, lane_id, seg, index) + assert lane_id not in self._lanes + self._lanes[lane_id] = lane + processed_ids.add(seg_id) + lanes.append(lane) + + # Create the road and fill in references + road = ArgoverseMap.Road(road_id, lanes) + assert road_id not in self._roads + self._roads[road_id] = road + for lane in lanes: + lane._road = road + + # Create lanes for the remaining lane segments, each with their own road + remaining_ids = all_ids - processed_ids + for seg_id in remaining_ids: + lane_seg = self._avm.vector_lane_segments[seg_id] + road_id = f"road-{lane_seg.id}" + lane_id = f"lane-{lane_seg.id}" + + lane = ArgoverseMap.Lane(self, lane_id, lane_seg, 0) + road = ArgoverseMap.Road(road_id, [lane]) + lane._road = road + + assert road_id not in self._roads + assert lane_id not in self._lanes + self._roads[road_id] = road + self._lanes[lane_id] = lane + + # Patch in incoming/outgoing lanes now that all lanes have been created + for lane in self._lanes.values(): + lane._incoming_lanes = [ + self.lane_by_id(f"lane-{seg_id}") + for seg_id in lane.lane_seg.predecessors + if seg_id in all_ids + ] + lane._outgoing_lanes = [ + self.lane_by_id(f"lane-{seg_id}") + for seg_id in lane.lane_seg.successors + if seg_id in all_ids + ] + + self._compute_lane_intersections() + + end = time.time() + elapsed = round((end - start) * 1000.0, 3) + self._log.info(f"Loading Argoverse map took: {elapsed} ms") + + @property + def source(self) -> str: + """Path to the directory containing the map JSON file.""" + return self._map_spec.source + + @cached_property + def bounding_box(self) -> Optional[BoundingBox]: + xs, ys = np.array([]), np.array([]) + for lane_seg in self._avm.get_scenario_lane_segments(): + xs = np.concatenate((xs, lane_seg.polygon_boundary[:, 0])) + ys = np.concatenate((ys, lane_seg.polygon_boundary[:, 1])) + + return BoundingBox( + min_pt=Point(x=np.min(xs), y=np.min(ys)), + max_pt=Point(x=np.max(xs), y=np.max(ys)), + ) + + def is_same_map(self, map_spec) -> bool: + return map_spec.source == self._map_spec.source + + def _compute_traffic_dividers(self) -> Tuple[List, List]: + lane_dividers = [] # divider between lanes with same traffic direction + road_dividers = [] # divider between roads with opposite traffic direction + processed_ids = [] + + for lane_seg in self._avm.get_scenario_lane_segments(): + if lane_seg.id in processed_ids: + continue + if lane_seg.right_neighbor_id is None: + cur_seg = lane_seg + while True: + if ( + cur_seg.left_neighbor_id is None + or cur_seg.id in processed_ids + or ( + cur_seg.left_neighbor_id + not in self._avm.vector_lane_segments + ) + ): + break # This is the leftmost lane in the road, so stop + else: + left_mark = cur_seg.left_lane_marking.mark_type + lane = self.lane_by_id(f"lane-{cur_seg.id}") + left_boundary = [(p[0], p[1]) for p in lane.left_pts] + if left_mark in ArgoverseMap.LANE_MARKINGS: + lane_dividers.append(left_boundary) + elif left_mark in ArgoverseMap.ROAD_MARKINGS: + road_dividers.append(left_boundary) + processed_ids.append(cur_seg.id) + cur_seg = self._avm.vector_lane_segments[ + cur_seg.left_neighbor_id + ] + + return lane_dividers, road_dividers + + def to_glb(self, glb_dir): + polygons = [] + for lane_id, lane in self._lanes.items(): + metadata = { + "road_id": lane.road.road_id, + "lane_id": lane_id, + "lane_index": lane.index, + } + polygons.append((lane.shape(), metadata)) + + lane_dividers, edge_dividers = self._compute_traffic_dividers() + + map_glb = make_map_glb( + polygons, self.bounding_box, lane_dividers, edge_dividers + ) + map_glb.write_glb(Path(glb_dir) / "map.glb") + + road_lines_glb = make_road_line_glb(edge_dividers) + road_lines_glb.write_glb(Path(glb_dir) / "road_lines.glb") + + lane_lines_glb = make_road_line_glb(lane_dividers) + lane_lines_glb.write_glb(Path(glb_dir) / "lane_lines.glb") + + class Surface(RoadMapWithCaches.Surface): + """Surface representation for Argoverse maps.""" + + def __init__(self, surface_id: str, road_map): + self._surface_id = surface_id + + @property + def surface_id(self) -> str: + return self._surface_id + + @property + def is_drivable(self) -> bool: + return True + + def surface_by_id(self, surface_id: str) -> RoadMap.Surface: + return self._surfaces.get(surface_id) + + class Lane(RoadMapWithCaches.Lane, Surface): + """Lane representation for Argoverse maps.""" + + def __init__( + self, map: "ArgoverseMap", lane_id: str, lane_seg: LaneSegment, index: int + ): + super().__init__(lane_id, map) + self._map = map + self._lane_id = lane_id + self.lane_seg = lane_seg + self._index = index + self._road = None + self._incoming_lanes = None + self._outgoing_lanes = None + self._intersections = None + + self._polygon = lane_seg.polygon_boundary[:, :2] + self._centerline = self._map._avm.get_lane_segment_centerline(lane_seg.id)[ + :, :2 + ] + + xs = self._polygon[:, 0] + ys = self._polygon[:, 1] + self._bbox = BoundingBox( + min_pt=Point(x=np.amin(xs), y=np.amin(ys)), + max_pt=Point(x=np.amax(xs), y=np.amax(ys)), + ) + + # Compute equally-spaced points for lane boundaries by interpolating + n = len(self._centerline) + self.left_pts = interp_arc( + n, points=self.lane_seg.left_lane_boundary.xyz[:, :2] + ) + self.right_pts = interp_arc( + n, points=self.lane_seg.right_lane_boundary.xyz[:, :2] + ) + + def __hash__(self) -> int: + return hash(self.lane_id) + + @property + def lane_id(self) -> str: + return self._lane_id + + @property + def road(self) -> RoadMap.Road: + return self._road + + @property + def speed_limit(self) -> Optional[float]: + return ArgoverseMap.DEFAULT_LANE_SPEED + + @lru_cache(maxsize=1024) + def width_at_offset(self, lane_point_s: float) -> Tuple[float, float]: + world_point = self.from_lane_coord( + RefLinePoint(lane_point_s, 0) + ).as_np_array[:2] + deltas = self._centerline - world_point + dists = np.linalg.norm(deltas, axis=1) + closest_index = np.argmin(dists) + p1 = self.left_pts[closest_index] + p2 = self.right_pts[closest_index] + width = np.linalg.norm(p2 - p1) + return width, 1.0 + + @cached_property + def length(self) -> float: + length = 0 + for p1, p2 in zip(self._centerline, self._centerline[1:]): + length += np.linalg.norm(p2 - p1) + return length + + @cached_property + def center_polyline(self) -> List[Point]: + return [Point(p[0], p[1]) for p in self._centerline] + + @property + def in_junction(self) -> bool: + return self.lane_seg.is_intersection + + @property + def index(self) -> int: + return self._index + + @lru_cache(maxsize=4) + def shape( + self, buffer_width: float = 0.0, default_width: Optional[float] = None + ) -> Polygon: + return Polygon(self._polygon) + + @cached_property + def lanes_in_same_direction(self) -> List[RoadMap.Lane]: + return [lane for lane in self.road.lanes if lane.lane_id != self.lane_id] + + @cached_property + def lane_to_left(self) -> Tuple[RoadMap.Lane, bool]: + result = None + for other in self.lanes_in_same_direction: + if other.index > self.index and ( + not result or other.index < result.index + ): + result = other + return result, True + + @cached_property + def lane_to_right(self) -> Tuple[RoadMap.Lane, bool]: + result = None + for other in self.lanes_in_same_direction: + if other.index < self.index and ( + not result or other.index > result.index + ): + result = other + return result, True + + @property + def incoming_lanes(self) -> List[RoadMap.Lane]: + return self._incoming_lanes + + @property + def outgoing_lanes(self) -> List[RoadMap.Lane]: + return self._outgoing_lanes + + @lru_cache(maxsize=16) + def oncoming_lanes_at_offset(self, offset: float) -> List[RoadMap.Lane]: + result = [] + radius = 1.1 * self.width_at_offset(offset)[0] + pt = self.from_lane_coord(RefLinePoint(offset)) + nearby_lanes = self._map.nearest_lanes(pt, radius=radius) + if not nearby_lanes: + return result + my_vect = self.vector_at_offset(offset) + my_norm = np.linalg.norm(my_vect) + if my_norm == 0: + return result + threshold = -0.995562 # cos(175*pi/180) + for lane, _ in nearby_lanes: + if lane == self: + continue + lane_refline_pt = lane.to_lane_coord(pt) + lv = lane.vector_at_offset(lane_refline_pt.s) + lv_norm = np.linalg.norm(lv) + if lv_norm == 0: + continue + lane_angle = np.dot(my_vect, lv) / (my_norm * lv_norm) + if lane_angle < threshold: + result.append(lane) + return result + + @cached_property + def foes(self) -> List[RoadMap.Lane]: + foes = set(self._intersections) + foes |= { + incoming + for outgoing in self.outgoing_lanes + for incoming in outgoing.incoming_lanes + if incoming != self + } + return list(foes) + + @lru_cache(maxsize=8) + def contains_point(self, point: Point) -> bool: + assert type(point) == Point + if ( + self._bbox.min_pt.x <= point[0] <= self._bbox.max_pt.x + and self._bbox.min_pt.y <= point[1] <= self._bbox.max_pt.y + ): + return self.shape().contains(point.as_shapely) + return False + + def waypoint_paths_for_pose( + self, pose: Pose, lookahead: int, route: RoadMap.Route = None + ) -> List[List[Waypoint]]: + if not self.is_drivable: + return [] + road_ids = [road.road_id for road in route.roads] if route else None + return self._waypoint_paths_at(pose.point, lookahead, road_ids) + + def waypoint_paths_at_offset( + self, offset: float, lookahead: int = 30, route: RoadMap.Route = None + ) -> List[List[Waypoint]]: + if not self.is_drivable: + return [] + wp_start = self.from_lane_coord(RefLinePoint(offset)) + road_ids = [road.road_id for road in route.roads] if route else None + return self._waypoint_paths_at(wp_start, lookahead, road_ids) + + def _waypoint_paths_at( + self, + point: Point, + lookahead: int, + filter_road_ids: Optional[Sequence[str]] = None, + ) -> List[List[Waypoint]]: + if not self.is_drivable: + return [] + closest_linked_lp = ( + self._map._lanepoints.closest_linked_lanepoint_on_lane_to_point( + point, self._lane_id + ) + ) + return self._map._waypoints_starting_at_lanepoint( + closest_linked_lp, + lookahead, + tuple(filter_road_ids) if filter_road_ids else (), + point, + ) + + def lane_by_id(self, lane_id: str) -> RoadMap.Lane: + lane = self._lanes.get(lane_id) + assert lane, f"ArgoverseMap got request for unknown lane_id: '{lane_id}'" + return lane + + def _build_lane_r_tree(self): + result = rtree.index.Index() + result.interleaved = True + for idx, lane in enumerate(self._lanes.values()): + xs = lane._polygon[:, 0] + ys = lane._polygon[:, 1] + bounding_box = ( + np.amin(xs), + np.amin(ys), + np.amax(xs), + np.amax(ys), + ) + result.add(idx, bounding_box) + return result + + def _get_neighboring_lanes( + self, x: float, y: float, r: float = 0.1 + ) -> List[Tuple[RoadMapWithCaches.Lane, float]]: + neighboring_lanes = [] + if self._lane_rtree is None: + self._lane_rtree = self._build_lane_r_tree() + + spt = SPoint(x, y) + lanes = list(self._lanes.values()) + for i in self._lane_rtree.intersection((x - r, y - r, x + r, y + r)): + lane = lanes[i] + d = lane.shape().distance(spt) + if d < r: + neighboring_lanes.append((lane, d)) + return neighboring_lanes + + @lru_cache(maxsize=1024) + def nearest_lanes( + self, + point: Point, + radius: Optional[float] = None, + include_junctions: bool = False, + ) -> List[Tuple[RoadMapWithCaches.Lane, float]]: + if radius is None: + radius = 5 + candidate_lanes = self._get_neighboring_lanes(point[0], point[1], r=radius) + candidate_lanes.sort(key=lambda lane_dist_tup: lane_dist_tup[1]) + return candidate_lanes + + @lru_cache(maxsize=16) + def road_with_point(self, point: Point) -> RoadMap.Road: + radius = 5 + for nl, dist in self.nearest_lanes(point, radius): + if nl.contains_point(point): + return nl.road + return None + + class Road(RoadMapWithCaches.Road, Surface): + """Road representation for Argoverse maps.""" + + def __init__(self, road_id: str, lanes: List[RoadMap.Lane]): + super().__init__(road_id, None) + self._road_id = road_id + self._lanes = lanes + + x_mins, y_mins, x_maxs, y_maxs = [], [], [], [] + for lane in self._lanes: + x_mins.append(lane._bbox.min_pt.x) + y_mins.append(lane._bbox.min_pt.y) + x_maxs.append(lane._bbox.max_pt.x) + y_maxs.append(lane._bbox.max_pt.y) + + self._bbox = BoundingBox( + min_pt=Point(x=min(x_mins), y=min(y_mins)), + max_pt=Point(x=max(x_maxs), y=max(y_maxs)), + ) + + def __hash__(self) -> int: + return hash(self.road_id) + + @property + def road_id(self) -> str: + return self._road_id + + @property + def composite_road(self) -> RoadMap.Road: + """Return an abstract Road composed of one or more RoadMap.Road segments + (including this one) that has been inferred to correspond to one continuous + real-world road. May return same object as self.""" + return self + + @property + def is_composite(self) -> bool: + """Returns True if this Road object was inferred + and composed out of subordinate Road objects.""" + return False + + @cached_property + def is_junction(self) -> bool: + for lane in self.lanes: + if lane.foes or len(lane.incoming_lanes) > 1: + return True + return False + + @cached_property + def length(self) -> float: + # Neighbouring lanes in Argoverse can be different lengths. Since this is + # just used for routes, we take the average lane length in this road. + return sum([lane.length for lane in self.lanes]) / len(self.lanes) + + @property + def incoming_roads(self) -> List[RoadMap.Road]: + return list( + {in_lane.road for lane in self.lanes for in_lane in lane.incoming_lanes} + ) + + @property + def outgoing_roads(self) -> List[RoadMap.Road]: + return list( + { + out_lane.road + for lane in self.lanes + for out_lane in lane.outgoing_lanes + } + ) + + @lru_cache(maxsize=16) + def oncoming_roads_at_point(self, point: Point) -> List[RoadMap.Road]: + result = [] + for lane in self.lanes: + offset = lane.to_lane_coord(point).s + result += [ + ol.road + for ol in lane.oncoming_lanes_at_offset(offset) + if ol.road != self + ] + return result + + @property + def parallel_roads(self) -> List[RoadMap.Road]: + return [] + + @property + def lanes(self) -> List[RoadMap.Lane]: + return self._lanes + + def lane_at_index(self, index: int) -> RoadMap.Lane: + return self.lanes[index] + + @lru_cache(maxsize=8) + def contains_point(self, point: Point) -> bool: + if ( + self._bbox.min_pt.x <= point[0] <= self._bbox.max_pt.x + and self._bbox.min_pt.y <= point[1] <= self._bbox.max_pt.y + ): + for lane in self._lanes: + if lane.contains_point(point): + return True + return False + + def road_by_id(self, road_id: str) -> RoadMap.Road: + road = self._roads.get(road_id) + assert road, f"ArgoverseMap got request for unknown road_id: '{road_id}'" + return road + + class Route(RouteWithCache): + """Describes a route between Argoverse roads.""" + + def __init__(self, road_map): + super().__init__(road_map) + self._roads = [] + self._length = 0 + + @property + def roads(self) -> List[RoadMap.Road]: + return self._roads + + @property + def road_length(self) -> float: + return self._length + + def _add_road(self, road: RoadMap.Road): + """Add a road to this route.""" + self._length += road.length + self._roads.append(road) + + @cached_property + def geometry(self) -> Sequence[Sequence[Tuple[float, float]]]: + return [list(road.shape(0.0).exterior.coords) for road in self.roads] + + @staticmethod + def _shortest_route(start: RoadMap.Road, end: RoadMap.Road) -> List[RoadMap.Road]: + queue = [(start.length, start.road_id, start)] + came_from = dict() + came_from[start] = None + cost_so_far = dict() + cost_so_far[start] = start.length + current = None + + # Dijkstra’s Algorithm + while queue: + (_, _, current) = heapq.heappop(queue) + current: RoadMap.Road + if current == end: + break + for out_road in current.outgoing_roads: + new_cost = cost_so_far[current] + out_road.length + if out_road not in cost_so_far or new_cost < cost_so_far[out_road]: + cost_so_far[out_road] = new_cost + came_from[out_road] = current + heapq.heappush(queue, (new_cost, out_road.road_id, out_road)) + + # This means we couldn't find a valid route since the queue is empty + if current != end: + return [] + + # Reconstruct path + current = end + path = [] + while current != start: + path.append(current) + current = came_from[current] + path.append(start) + path.reverse() + return path + + def generate_routes( + self, + start_road: RoadMap.Road, + end_road: RoadMap.Road, + via: Optional[Sequence[RoadMap.Road]] = None, + max_to_gen: int = 1, + ) -> List[RoadMap.Route]: + assert ( + max_to_gen == 1 + ), "multiple route generation not yet supported for Argoverse" + new_route = ArgoverseMap.Route(self) + result = [new_route] + + roads = [start_road] + if via: + roads += via + if end_road != start_road: + roads.append(end_road) + + route_roads = [] + for cur_road, next_road in zip(roads, roads[1:] + [None]): + if not next_road: + route_roads.append(cur_road) + break + sub_route = ArgoverseMap._shortest_route(cur_road, next_road) or [] + if len(sub_route) < 2: + self._log.warning( + f"Unable to find valid path between {(cur_road.road_id, next_road.road_id)}." + ) + return result + # The sub route includes the boundary roads (cur_road, next_road). + # We clip the latter to prevent duplicates + route_roads.extend(sub_route[:-1]) + + for road in route_roads: + new_route._add_road(road) + return result + + def random_route( + self, + max_route_len: int = 10, + starting_road: Optional[RoadMap.Road] = None, + only_drivable: bool = True, + ) -> RoadMap.Route: + assert not starting_road or not only_drivable or starting_road.is_drivable + route = ArgoverseMap.Route(self) + next_roads = [starting_road] if starting_road else list(self._roads.values()) + if only_drivable: + next_roads = [r for r in next_roads if r.is_drivable] + while next_roads and len(route.roads) < max_route_len: + cur_road = random.choice(next_roads) + route._add_road(cur_road) + next_roads = list(cur_road.outgoing_roads) + return route + + def empty_route(self) -> RoadMap.Route: + return ArgoverseMap.Route(self) + + def route_from_road_ids(self, road_ids: Sequence[str]) -> RoadMap.Route: + return ArgoverseMap.Route.from_road_ids(self, road_ids) + + class _WaypointsCache: + def __init__(self): + self.lookahead = 0 + self.point = Point(0, 0) + self.filter_road_ids = () + self._starts = {} + + # XXX: all vehicles share this cache now (as opposed to before + # when it was in Plan.py and each vehicle had its own cache). + # TODO: probably need to add vehicle_id to the key somehow (or just make it bigger) + def _match(self, lookahead, point, filter_road_ids) -> bool: + return ( + lookahead <= self.lookahead + and point[0] == self.point[0] + and point[1] == self.point[1] + and filter_road_ids == self.filter_road_ids + ) + + def update( + self, + lookahead: int, + point: Point, + filter_road_ids: tuple, + llp, + paths: List[List[Waypoint]], + ): + """Update the current cache if not already cached.""" + if not self._match(lookahead, point, filter_road_ids): + self.lookahead = lookahead + self.point = point + self.filter_road_ids = filter_road_ids + self._starts = {} + self._starts[llp.lp.lane.index] = paths + + def query( + self, + lookahead: int, + point: Point, + filter_road_ids: tuple, + llp, + ) -> Optional[List[List[Waypoint]]]: + """Attempt to find previously cached waypoints""" + if self._match(lookahead, point, filter_road_ids): + hit = self._starts.get(llp.lp.lane.index, None) + if hit: + # consider just returning all of them (not slicing)? + return [path[: (lookahead + 1)] for path in hit] + return None + + def waypoint_paths( + self, + pose: Pose, + lookahead: int, + within_radius: float = 5, + route: RoadMap.Route = None, + ) -> List[List[Waypoint]]: + road_ids = [] + if route and route.roads: + road_ids = [road.road_id for road in route.roads] + if road_ids: + return self._waypoint_paths_along_route(pose.point, lookahead, road_ids) + closest_lps = self._lanepoints.closest_lanepoints( + [pose], within_radius=within_radius + ) + closest_lane = closest_lps[0].lane + waypoint_paths = [] + for lane in closest_lane.road.lanes: + waypoint_paths += lane._waypoint_paths_at(pose.point, lookahead) + return sorted(waypoint_paths, key=lambda p: p[0].lane_index) + + def _waypoint_paths_along_route( + self, point: Point, lookahead: int, route: Sequence[str] + ) -> List[List[Waypoint]]: + """finds the closest lane to vehicle's position that is on its route, + then gets waypoint paths from all lanes in its road there.""" + assert len(route) > 0, f"Expected at least 1 road in the route, got: {route}" + closest_llp_on_each_route_road = [ + self._lanepoints.closest_linked_lanepoint_on_road(point, road) + for road in route + ] + closest_linked_lp = min( + closest_llp_on_each_route_road, + key=lambda l_lp: np.linalg.norm( + vec_2d(l_lp.lp.pose.position) - vec_2d(point) + ), + ) + closest_lane = closest_linked_lp.lp.lane + waypoint_paths = [] + for lane in closest_lane.road.lanes: + waypoint_paths += lane._waypoint_paths_at(point, lookahead, route) + + return sorted(waypoint_paths, key=len, reverse=True) + + @staticmethod + def _equally_spaced_path( + path: Sequence[LinkedLanePoint], + point: Point, + lp_spacing: float, + ) -> List[Waypoint]: + """given a list of LanePoints starting near point, return corresponding + Waypoints that may not be evenly spaced (due to lane change) but start at point. + """ + + continuous_variables = [ + "positions_x", + "positions_y", + "headings", + "lane_width", + "speed_limit", + "lane_offset", + ] + discrete_variables = ["lane_id", "lane_index"] + + ref_lanepoints_coordinates = { + parameter: [] for parameter in (continuous_variables + discrete_variables) + } + for idx, lanepoint in enumerate(path): + if lanepoint.is_inferred and 0 < idx < len(path) - 1: + continue + + ref_lanepoints_coordinates["positions_x"].append( + lanepoint.lp.pose.position[0] + ) + ref_lanepoints_coordinates["positions_y"].append( + lanepoint.lp.pose.position[1] + ) + ref_lanepoints_coordinates["headings"].append( + lanepoint.lp.pose.heading.as_bullet + ) + ref_lanepoints_coordinates["lane_id"].append(lanepoint.lp.lane.lane_id) + ref_lanepoints_coordinates["lane_index"].append(lanepoint.lp.lane.index) + + ref_lanepoints_coordinates["lane_width"].append(lanepoint.lp.lane_width) + + ref_lanepoints_coordinates["lane_offset"].append( + lanepoint.lp.lane.offset_along_lane(lanepoint.lp.pose.point) + ) + + ref_lanepoints_coordinates["speed_limit"].append( + lanepoint.lp.lane.speed_limit + ) + + ref_lanepoints_coordinates["headings"] = inplace_unwrap( + ref_lanepoints_coordinates["headings"] + ) + first_lp_heading = ref_lanepoints_coordinates["headings"][0] + lp_position = path[0].lp.pose.point.as_np_array[:2] + vehicle_pos = point.as_np_array[:2] + heading_vec = radians_to_vec(first_lp_heading) + projected_distant_lp_vehicle = np.inner( + (vehicle_pos - lp_position), heading_vec + ) + + ref_lanepoints_coordinates["positions_x"][0] = ( + lp_position[0] + projected_distant_lp_vehicle * heading_vec[0] + ) + ref_lanepoints_coordinates["positions_y"][0] = ( + lp_position[1] + projected_distant_lp_vehicle * heading_vec[1] + ) + + cumulative_path_dist = np.cumsum( + np.sqrt( + np.ediff1d(ref_lanepoints_coordinates["positions_x"], to_begin=0) ** 2 + + np.ediff1d(ref_lanepoints_coordinates["positions_y"], to_begin=0) ** 2 + ) + ) + + if len(cumulative_path_dist) <= lp_spacing: + lp = path[0].lp + + return [ + Waypoint( + pos=lp.pose.position[:2], + heading=lp.pose.heading, + lane_width=lp.lane.width_at_offset(0)[0], + speed_limit=lp.lane.speed_limit, + lane_id=lp.lane.lane_id, + lane_index=lp.lane.index, + lane_offset=lp.lane.offset_along_lane(lp.pose.point), + ) + ] + + evenly_spaced_cumulative_path_dist = np.linspace( + 0, cumulative_path_dist[-1], len(path) + ) + + evenly_spaced_coordinates = {} + for variable in continuous_variables: + evenly_spaced_coordinates[variable] = np.interp( + evenly_spaced_cumulative_path_dist, + cumulative_path_dist, + ref_lanepoints_coordinates[variable], + ) + + for variable in discrete_variables: + ref_coordinates = ref_lanepoints_coordinates[variable] + evenly_spaced_coordinates[variable] = [] + jdx = 0 + for idx in range(len(path)): + while ( + jdx + 1 < len(cumulative_path_dist) + and evenly_spaced_cumulative_path_dist[idx] + > cumulative_path_dist[jdx + 1] + ): + jdx += 1 + + evenly_spaced_coordinates[variable].append(ref_coordinates[jdx]) + evenly_spaced_coordinates[variable].append(ref_coordinates[-1]) + + waypoint_path = [] + for idx in range(len(path)): + waypoint_path.append( + Waypoint( + pos=np.array( + [ + evenly_spaced_coordinates["positions_x"][idx], + evenly_spaced_coordinates["positions_y"][idx], + ] + ), + heading=Heading(evenly_spaced_coordinates["headings"][idx]), + lane_width=evenly_spaced_coordinates["lane_width"][idx], + speed_limit=evenly_spaced_coordinates["speed_limit"][idx], + lane_id=evenly_spaced_coordinates["lane_id"][idx], + lane_index=evenly_spaced_coordinates["lane_index"][idx], + lane_offset=evenly_spaced_coordinates["lane_offset"][idx], + ) + ) + + return waypoint_path + + def _waypoints_starting_at_lanepoint( + self, + lanepoint: LinkedLanePoint, + lookahead: int, + filter_road_ids: tuple, + point: Point, + ) -> List[List[Waypoint]]: + """computes equally-spaced Waypoints for all lane paths starting at lanepoint + up to lookahead waypoints ahead, constrained to filter_road_ids if specified.""" + + # The following acts sort of like lru_cache(1), but it allows + # for lookahead to be <= to the cached value... + cache_paths = self._waypoints_cache.query( + lookahead, point, filter_road_ids, lanepoint + ) + if cache_paths: + return cache_paths + + lanepoint_paths = self._lanepoints.paths_starting_at_lanepoint( + lanepoint, lookahead, filter_road_ids + ) + result = [ + ArgoverseMap._equally_spaced_path( + path, + point, + self._map_spec.lanepoint_spacing, + ) + for path in lanepoint_paths + ] + + self._waypoints_cache.update( + lookahead, point, filter_road_ids, lanepoint, result + ) + + return result diff --git a/smarts/core/default_map_builder.py b/smarts/core/default_map_builder.py index 20f785ff9e..dba71d6aa3 100644 --- a/smarts/core/default_map_builder.py +++ b/smarts/core/default_map_builder.py @@ -55,6 +55,7 @@ def _clear_cache(): _SUMO_MAP = 1 _OPENDRIVE_MAP = 2 _WAYMO_MAP = 3 +_ARGOVERSE_MAP = 4 def find_mapfile_in_dir(map_dir: str) -> Tuple[int, str]: @@ -78,6 +79,9 @@ def find_mapfile_in_dir(map_dir: str) -> Tuple[int, str]: elif ".tfrecord" in f: map_type = _WAYMO_MAP map_path = os.path.join(map_dir, f) + elif "log_map_archive" in f: + map_type = _ARGOVERSE_MAP + map_path = os.path.join(map_dir, f) return map_type, map_path @@ -132,17 +136,25 @@ def get_road_map(map_spec) -> Tuple[Optional[RoadMap], Optional[str]]: except (ImportError, ModuleNotFoundError): print(sys.exc_info()) print( - "You may not have installed the [waymo] dependencies required to build and use WaymoMap Scenario. Install them first using the command `pip install -e .[waymo]` at the source directory." + "You may not have installed the [waymo] dependencies required to build and use WaymoMap scenarios. Install them first using the command `pip install -e .[waymo]` at the source directory." ) return None, None map_class = WaymoMap + elif map_type == _ARGOVERSE_MAP: + try: + from smarts.core.argoverse_map import ArgoverseMap # pytype: disable=import-error + except (ImportError, ModuleNotFoundError): + print(sys.exc_info()) + print( + "You may not have installed the [argoverse] dependencies required to build and use Argoverse scenarios. Install them first using the command `pip install -e .[argoverse]` at the source directory." + ) + return None, None + map_class = ArgoverseMap else: return None, None if _existing_map: - if isinstance(_existing_map.obj, map_class) and _existing_map.obj.is_same_map( - map_spec - ): + if isinstance(_existing_map.obj, map_class) and _existing_map.obj.is_same_map(map_spec): return _existing_map.obj, _existing_map.map_hash _clear_cache() diff --git a/smarts/core/lanepoints.py b/smarts/core/lanepoints.py index 3793285d4d..0527b6894d 100644 --- a/smarts/core/lanepoints.py +++ b/smarts/core/lanepoints.py @@ -491,6 +491,125 @@ def _shape_lanepoints_along_lane( return cls(shape_lps, spacing) + @classmethod + def from_argoverse( + cls, + argoverse_map, + spacing, + ): + """Computes the lane shape (start/shape/end) lanepoints for all lanes in + the network, the result of this function can be used to interpolate + lanepoints along lanes to the desired granularity. + """ + from .argoverse_map import ArgoverseMap + + assert type(argoverse_map) == ArgoverseMap + + def _shape_lanepoints_along_lane( + lane: RoadMap.Lane, + lanepoint_by_lane_memo: dict, + ) -> Tuple[LinkedLanePoint, List[LinkedLanePoint]]: + lane_queue = queue.Queue() + lane_queue.put((lane, None)) + shape_lanepoints = [] + initial_lanepoint = None + while not lane_queue.empty(): + curr_lane, previous_lp = lane_queue.get() + first_lanepoint = lanepoint_by_lane_memo.get(curr_lane.lane_id) + if first_lanepoint: + if previous_lp: + previous_lp.nexts.append(first_lanepoint) + continue + + lane_shape = curr_lane._centerline + + assert ( + len(lane_shape) >= 2 + ), f"{repr(lane_shape)} for lane_id={curr_lane.lane_id}" + + vd = lane_shape[1] - lane_shape[0] + heading = Heading(vec_to_radians(vd[:2])) + orientation = fast_quaternion_from_angle(heading) + + lane_width, _ = curr_lane.width_at_offset(0) + first_lanepoint = LinkedLanePoint( + lp=LanePoint( + lane=curr_lane, + pose=Pose(position=lane_shape[0], orientation=orientation), + lane_width=lane_width, + ), + nexts=[], + is_inferred=False, + ) + + if previous_lp is not None: + previous_lp.nexts.append(first_lanepoint) + + if initial_lanepoint is None: + initial_lanepoint = first_lanepoint + + lanepoint_by_lane_memo[curr_lane.lane_id] = first_lanepoint + shape_lanepoints.append(first_lanepoint) + curr_lanepoint = first_lanepoint + + for p1, p2 in zip(lane_shape[1:], lane_shape[2:]): + vd = p2 - p1 + heading_ = Heading(vec_to_radians(vd[:2])) + orientation_ = fast_quaternion_from_angle(heading_) + lane_width, _ = curr_lane.width_at_offset(0) + linked_lanepoint = LinkedLanePoint( + lp=LanePoint( + lane=curr_lane, + pose=Pose(position=p1, orientation=orientation_), + lane_width=lane_width, + ), + nexts=[], + is_inferred=False, + ) + + shape_lanepoints.append(linked_lanepoint) + curr_lanepoint.nexts.append(linked_lanepoint) + curr_lanepoint = linked_lanepoint + + # Add a lanepoint for the last point of the current lane + curr_lanepoint_lane = curr_lanepoint.lp.lane + lane_width, _ = curr_lanepoint_lane.width_at_offset(0) + last_linked_lanepoint = LinkedLanePoint( + lp=LanePoint( + lane=curr_lanepoint.lp.lane, + pose=Pose( + position=lane_shape[-1], + orientation=curr_lanepoint.lp.pose.orientation, + ), + lane_width=lane_width, + ), + nexts=[], + is_inferred=False, + ) + + shape_lanepoints.append(last_linked_lanepoint) + curr_lanepoint.nexts.append(last_linked_lanepoint) + curr_lanepoint = last_linked_lanepoint + + for out_lane in curr_lane.outgoing_lanes: + if out_lane and out_lane.is_drivable: + lane_queue.put((out_lane, curr_lanepoint)) + return initial_lanepoint, shape_lanepoints + + roads = argoverse_map._roads + lanepoint_by_lane_memo = {} + shape_lps = [] + + for road in roads.values(): + for lane in road.lanes: + if lane.is_drivable: + _, new_lps = _shape_lanepoints_along_lane( + lane, lanepoint_by_lane_memo + ) + shape_lps += new_lps + + return cls(shape_lps, spacing) + @staticmethod def _build_kd_tree(linked_lps: Sequence[LinkedLanePoint]) -> KDTree: return KDTree( diff --git a/smarts/core/opendrive_road_network.py b/smarts/core/opendrive_road_network.py index 9824ccae92..c9c9a0b8b7 100644 --- a/smarts/core/opendrive_road_network.py +++ b/smarts/core/opendrive_road_network.py @@ -32,10 +32,10 @@ import numpy as np import rtree -import trimesh -import trimesh.scene from cached_property import cached_property +from smarts.core.utils.glb import make_map_glb, make_road_line_glb + # pytype: disable=import-error @@ -75,11 +75,9 @@ from shapely.geometry import Point as SPoint from shapely.geometry import Polygon -from trimesh.exchange import gltf from smarts.core.road_map import RoadMap, RoadMapWithCaches, Waypoint from smarts.core.route_cache import RouteWithCache -from smarts.core.utils.geometry import generate_meshes_from_polygons from smarts.core.utils.key_wrapper import KeyWrapper from smarts.core.utils.math import ( CubicPolynomial, @@ -96,35 +94,6 @@ from .lanepoints import LanePoints, LinkedLanePoint -def _convert_camera(camera): - result = { - "name": camera.name, - "type": "perspective", - "perspective": { - "aspectRatio": camera.fov[0] / camera.fov[1], - "yfov": np.radians(camera.fov[1]), - "znear": float(camera.z_near), - # HACK: The trimesh gltf export doesn't include a zfar which Panda3D GLB - # loader expects. Here we override to make loading possible. - "zfar": float(camera.z_near + 100), - }, - } - return result - - -gltf._convert_camera = _convert_camera - - -class _GLBData: - def __init__(self, bytes_): - self._bytes = bytes_ - - def write_glb(self, output_path: str): - """Generate a geometry file.""" - with open(output_path, "wb") as f: - f.write(self._bytes) - - @dataclass class LaneBoundary: """Describes a lane boundary.""" @@ -763,32 +732,8 @@ def bounding_box(self) -> BoundingBox: ) def to_glb(self, glb_dir): - lane_dividers, edge_dividers = self._compute_traffic_dividers() - map_glb = self._make_glb_from_polys(lane_dividers, edge_dividers) - map_glb.write_glb(Path(glb_dir) / "map.glb") - - road_lines_glb = self._make_road_line_glb(edge_dividers) - road_lines_glb.write_glb(Path(glb_dir) / "road_lines.glb") - - lane_lines_glb = self._make_road_line_glb(lane_dividers) - lane_lines_glb.write_glb(Path(glb_dir) / "lane_lines.glb") - - def _make_road_line_glb(self, lines: List[List[Tuple[float, float]]]): - scene = trimesh.Scene() - for line_pts in lines: - vertices = [(*pt, 0.1) for pt in line_pts] - point_cloud = trimesh.PointCloud(vertices=vertices) - point_cloud.apply_transform( - trimesh.transformations.rotation_matrix(math.pi / 2, [-1, 0, 0]) - ) - scene.add_geometry(point_cloud) - return _GLBData(gltf.export_glb(scene)) - - def _make_glb_from_polys(self, lane_dividers, edge_dividers): - scene = trimesh.Scene() polygons = [] - for lane_id in self._lanes: - lane = self._lanes[lane_id] + for lane_id, lane in self._lanes.items(): metadata = { "road_id": lane.road.road_id, "lane_id": lane_id, @@ -796,36 +741,18 @@ def _make_glb_from_polys(self, lane_dividers, edge_dividers): } polygons.append((lane.shape(), metadata)) - meshes = generate_meshes_from_polygons(polygons) - - # Attach additional information for rendering as metadata in the map glb - # <2D-BOUNDING_BOX>: four floats separated by ',' (,,,), - # which describe x-minimum, y-minimum, x-maximum, and y-maximum - metadata = { - "bounding_box": ( - self.bounding_box.min_pt.x, - self.bounding_box.min_pt.y, - self.bounding_box.max_pt.x, - self.bounding_box.max_pt.y, - ) - } + lane_dividers, edge_dividers = self._compute_traffic_dividers() - # lane markings information - metadata["lane_dividers"] = lane_dividers - metadata["edge_dividers"] = edge_dividers + map_glb = make_map_glb( + polygons, self.bounding_box, lane_dividers, edge_dividers + ) + map_glb.write_glb(Path(glb_dir) / "map.glb") - for mesh in meshes: - mesh.visual = trimesh.visual.TextureVisuals( - material=trimesh.visual.material.PBRMaterial() - ) + road_lines_glb = make_road_line_glb(edge_dividers) + road_lines_glb.write_glb(Path(glb_dir) / "road_lines.glb") - road_id = mesh.metadata["road_id"] - lane_id = mesh.metadata.get("lane_id") - name = f"{road_id}" - if lane_id is not None: - name += f"-{lane_id}" - scene.add_geometry(mesh, name, extras=mesh.metadata) - return _GLBData(gltf.export_glb(scene, extras=metadata, include_normals=True)) + lane_lines_glb = make_road_line_glb(lane_dividers) + lane_lines_glb.write_glb(Path(glb_dir) / "lane_lines.glb") def _compute_traffic_dividers(self): lane_dividers = [] # divider between lanes with same traffic direction @@ -1215,7 +1142,9 @@ def center_at_point(self, point: Point) -> Point: return super().center_at_point(point) @lru_cache(8) - def _edges_at_point(self, point: Point) -> Tuple[Point, Point]: + def _edges_at_point( + self, point: Point + ) -> Tuple[Optional[Point], Optional[Point]]: """Get the boundary points perpendicular to the center of the lane closest to the given world coordinate. Args: @@ -1228,15 +1157,20 @@ def _edges_at_point(self, point: Point) -> Tuple[Point, Point]: reference_line_vertices_len = int((len(self._lane_polygon) - 1) / 2) # left_edge - left_edge_shape = self._lane_polygon[:reference_line_vertices_len] - left_offset = offset_along_shape(point[:2], left_edge_shape) + left_edge_shape = [ + Point(x, y) for x, y in self._lane_polygon[:reference_line_vertices_len] + ] + left_offset = offset_along_shape(point, left_edge_shape) left_edge = position_at_shape_offset(left_edge_shape, left_offset) # right_edge - right_edge_shape = self._lane_polygon[ - reference_line_vertices_len : len(self._lane_polygon) - 1 + right_edge_shape = [ + Point(x, y) + for x, y in self._lane_polygon[ + reference_line_vertices_len : len(self._lane_polygon) - 1 + ] ] - right_offset = offset_along_shape(point[:2], right_edge_shape) + right_offset = offset_along_shape(point, right_edge_shape) right_edge = position_at_shape_offset(right_edge_shape, right_offset) return left_edge, right_edge @@ -1766,7 +1700,8 @@ def _equally_spaced_path( width_threshold=None, ) -> List[Waypoint]: """given a list of LanePoints starting near point, return corresponding - Waypoints that may not be evenly spaced (due to lane change) but start at point.""" + Waypoints that may not be evenly spaced (due to lane change) but start at point. + """ continuous_variables = [ "positions_x", @@ -1785,7 +1720,6 @@ def _equally_spaced_path( skip_lanepoints = False index_skipped = [] for idx, lanepoint in enumerate(path): - if lanepoint.is_inferred and 0 < idx < len(path) - 1: continue @@ -1956,7 +1890,7 @@ def _waypoints_starting_at_lanepoint( lanepoint: LinkedLanePoint, lookahead: int, filter_road_ids: tuple, - point: Tuple[float, float, float], + point: Point, ) -> List[List[Waypoint]]: """computes equally-spaced Waypoints for all lane paths starting at lanepoint up to lookahead waypoints ahead, constrained to filter_road_ids if specified.""" diff --git a/smarts/core/road_map.py b/smarts/core/road_map.py index 37867c3d4d..6f2d168a26 100644 --- a/smarts/core/road_map.py +++ b/smarts/core/road_map.py @@ -233,7 +233,8 @@ def shape( """Returns a convex polygon representing this surface, buffered by buffered_width (which must be non-negative), where buffer_width is a buffer around the perimeter of the polygon. In some situations, it may be desirable to also specify a `default_width`, in which case the returned polygon should have a convex shape where the - distance across it is no less than buffered_width + default_width at any point.""" + distance across it is no less than buffered_width + default_width at any point. + """ raise NotImplementedError() def contains_point(self, point: Point) -> bool: @@ -295,7 +296,8 @@ def in_junction(self) -> bool: @property def index(self) -> int: """when not in_junction, 0 is outer / right-most (relative to lane heading) lane on road. - otherwise, index scheme is implementation-dependent, but must be deterministic.""" + otherwise, index scheme is implementation-dependent, but must be deterministic. + """ # TAI: UK roads raise NotImplementedError() @@ -697,7 +699,8 @@ def next_junction( ) -> Optional[Tuple[RoadMap.Lane, float]]: """Returns a lane within the next junction along the route from beginning of the current lane to the returned lane it connects with in the junction, - and the distance to it from this offset, or (None, inf) if there aren't any.""" + and the distance to it from this offset, or (None, inf) if there aren't any. + """ raise NotImplementedError() diff --git a/smarts/core/sumo_road_network.py b/smarts/core/sumo_road_network.py index 20262706c2..5322b1621c 100644 --- a/smarts/core/sumo_road_network.py +++ b/smarts/core/sumo_road_network.py @@ -18,7 +18,6 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN # THE SOFTWARE. import logging -import math import os import random from functools import lru_cache @@ -27,13 +26,10 @@ from typing import Any, Dict, List, Optional, Sequence, Set, Tuple import numpy as np -import trimesh -import trimesh.scene from cached_property import cached_property from shapely.geometry import Point as shPoint from shapely.geometry import Polygon from shapely.ops import nearest_points, snap -from trimesh.exchange import gltf from smarts.sstudio.types import MapSpec @@ -41,42 +37,14 @@ from .lanepoints import LanePoints, LinkedLanePoint from .road_map import RoadMap, Waypoint from .route_cache import RouteWithCache -from .utils.geometry import buffered_shape, generate_meshes_from_polygons +from .utils.geometry import buffered_shape +from .utils.glb import make_map_glb, make_road_line_glb from .utils.math import inplace_unwrap, radians_to_vec, vec_2d from smarts.core.utils.sumo import sumolib # isort:skip from sumolib.net.edge import Edge # isort:skip -def _convert_camera(camera): - result = { - "name": camera.name, - "type": "perspective", - "perspective": { - "aspectRatio": camera.fov[0] / camera.fov[1], - "yfov": np.radians(camera.fov[1]), - "znear": float(camera.z_near), - # HACK: The trimesh gltf export doesn't include a zfar which Panda3D GLB - # loader expects. Here we override to make loading possible. - "zfar": float(camera.z_near + 100), - }, - } - return result - - -gltf._convert_camera = _convert_camera - - -class _GLBData: - def __init__(self, bytes_): - self._bytes = bytes_ - - def write_glb(self, output_path): - """Generate a `.glb` geometry file.""" - with open(output_path, "wb") as f: - f.write(self._bytes) - - class SumoRoadNetwork(RoadMap): """A road network for a SUMO source.""" @@ -282,15 +250,15 @@ def scale_factor(self) -> float: return self._default_lane_width / SumoRoadNetwork.DEFAULT_LANE_WIDTH def to_glb(self, glb_dir): - lane_dividers, edge_dividers = self._compute_traffic_dividers() polys = self._compute_road_polygons() - map_glb = self._make_glb_from_polys(polys, lane_dividers, edge_dividers) + lane_dividers, edge_dividers = self._compute_traffic_dividers() + map_glb = make_map_glb(polys, self.bounding_box, lane_dividers, edge_dividers) map_glb.write_glb(Path(glb_dir) / "map.glb") - road_lines_glb = self._make_road_line_glb(edge_dividers) + road_lines_glb = make_road_line_glb(edge_dividers) road_lines_glb.write_glb(Path(glb_dir) / "road_lines.glb") - lane_lines_glb = self._make_road_line_glb(lane_dividers) + lane_lines_glb = make_road_line_glb(lane_dividers) lane_lines_glb.write_glb(Path(glb_dir) / "lane_lines.glb") class Surface(RoadMap.Surface): @@ -1279,44 +1247,6 @@ def _snap_external_holes(self, lane_to_poly, snap_threshold=2): if new_coords: lane_to_poly[lane_id] = (Polygon(new_coords), metadata) - def _make_road_line_glb(self, lines: List[List[Tuple[float, float]]]): - scene = trimesh.Scene() - for line_pts in lines: - vertices = [(*pt, 0.1) for pt in line_pts] - point_cloud = trimesh.PointCloud(vertices=vertices) - point_cloud.apply_transform( - trimesh.transformations.rotation_matrix(math.pi / 2, [-1, 0, 0]) - ) - scene.add_geometry(point_cloud) - return _GLBData(gltf.export_glb(scene)) - - def _make_glb_from_polys(self, polygons, lane_dividers, edge_dividers): - scene = trimesh.Scene() - meshes = generate_meshes_from_polygons(polygons) - # Attach additional information for rendering as metadata in the map glb - metadata = {} - - # <2D-BOUNDING_BOX>: four floats separated by ',' (,,,), - # which describe x-minimum, y-minimum, x-maximum, and y-maximum - metadata["bounding_box"] = self._graph.getBoundary() - - # lane markings information - metadata["lane_dividers"] = lane_dividers - metadata["edge_dividers"] = edge_dividers - - for mesh in meshes: - mesh.visual = trimesh.visual.TextureVisuals( - material=trimesh.visual.material.PBRMaterial() - ) - - road_id = mesh.metadata["road_id"] - lane_id = mesh.metadata.get("lane_id") - name = f"{road_id}" - if lane_id is not None: - name += f"-{lane_id}" - scene.add_geometry(mesh, name, extras=mesh.metadata) - return _GLBData(gltf.export_glb(scene, extras=metadata, include_normals=True)) - def _compute_traffic_dividers(self, threshold=1): lane_dividers = [] # divider between lanes with same traffic direction edge_dividers = [] # divider between lanes with opposite traffic direction @@ -1473,7 +1403,8 @@ def _equally_spaced_path( lp_spacing: float, ) -> List[Waypoint]: """given a list of LanePoints starting near point, that may not be evenly spaced, - returns the same number of Waypoints that are evenly spaced and start at point.""" + returns the same number of Waypoints that are evenly spaced and start at point. + """ continuous_variables = [ "positions_x", diff --git a/smarts/core/tests/test_argoverse.py b/smarts/core/tests/test_argoverse.py new file mode 100644 index 0000000000..3c59da96d6 --- /dev/null +++ b/smarts/core/tests/test_argoverse.py @@ -0,0 +1,49 @@ +# MIT License +# +# Copyright (C) 2022. Huawei Technologies Co., Ltd. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import subprocess +import sys + + +def test_argoverse_package_compatibility(): + """Test that the av2 package is installed and has no conflicts with the rest of the SMARTS dependencies.""" + + try: + import av2 + except ModuleNotFoundError: + print( + "The argoverse dependencies must be installed using `pip install -e .[argoverse]`" + ) + raise + + try: + subprocess.check_call( + [ + sys.executable, + "-m", + "pip", + "check", + ] + ) + except subprocess.CalledProcessError: + print("There is an incompatibility with the installed packages") + raise diff --git a/smarts/core/utils/geometry.py b/smarts/core/utils/geometry.py index b254b1a1d1..78cced13b1 100644 --- a/smarts/core/utils/geometry.py +++ b/smarts/core/utils/geometry.py @@ -18,11 +18,6 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN # THE SOFTWARE. -import math -from typing import Any, Dict, List, Tuple - -import numpy as np -import trimesh from shapely.geometry import LineString, MultiPolygon, Polygon from shapely.geometry.base import CAP_STYLE, JOIN_STYLE from shapely.ops import triangulate @@ -53,48 +48,3 @@ def triangulate_polygon(polygon: Polygon): for tri_face in triangulate(polygon) if tri_face.centroid.within(polygon) ] - - -def generate_meshes_from_polygons( - polygons: List[Tuple[Polygon, Dict[str, Any]]] -) -> List[trimesh.Trimesh]: - """Creates a mesh out of a list of polygons.""" - meshes = [] - - # Trimesh's API require a list of vertices and a list of faces, where each - # face contains three indexes into the vertices list. Ideally, the vertices - # are all unique and the faces list references the same indexes as needed. - # TODO: Batch the polygon processing. - for poly, metadata in polygons: - vertices, faces = [], [] - point_dict = dict() - current_point_index = 0 - - # Collect all the points on the shape to reduce checks by 3 times - for x, y in poly.exterior.coords: - p = (x, y, 0) - if p not in point_dict: - vertices.append(p) - point_dict[p] = current_point_index - current_point_index += 1 - triangles = triangulate_polygon(poly) - for triangle in triangles: - face = np.array( - [point_dict.get((x, y, 0), -1) for x, y in triangle.exterior.coords] - ) - # Add face if not invalid - if -1 not in face: - faces.append(face) - - if not vertices or not faces: - continue - - mesh = trimesh.Trimesh(vertices=vertices, faces=faces, metadata=metadata) - - # Trimesh doesn't support a coordinate-system="z-up" configuration, so we - # have to apply the transformation manually. - mesh.apply_transform( - trimesh.transformations.rotation_matrix(math.pi / 2, [-1, 0, 0]) - ) - meshes.append(mesh) - return meshes diff --git a/smarts/core/utils/glb.py b/smarts/core/utils/glb.py new file mode 100644 index 0000000000..af9917cf99 --- /dev/null +++ b/smarts/core/utils/glb.py @@ -0,0 +1,156 @@ +# Copyright (C) 2022. Huawei Technologies Co., Ltd. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +from pathlib import Path +from typing import Any, Dict, List, Tuple, Union + +import numpy as np +import trimesh +from shapely.geometry import Polygon +from trimesh.exchange import gltf + +from smarts.core.coordinates import BoundingBox +from smarts.core.utils.geometry import triangulate_polygon + + +def _convert_camera(camera): + result = { + "name": camera.name, + "type": "perspective", + "perspective": { + "aspectRatio": camera.fov[0] / camera.fov[1], + "yfov": np.radians(camera.fov[1]), + "znear": float(camera.z_near), + # HACK: The trimesh gltf export doesn't include a zfar which Panda3D GLB + # loader expects. Here we override to make loading possible. + "zfar": float(camera.z_near + 100), + }, + } + return result + + +gltf._convert_camera = _convert_camera + + +class GLBData: + """Convenience class for writing GLB files.""" + + def __init__(self, bytes_): + self._bytes = bytes_ + + def write_glb(self, output_path: Union[str, Path]): + """Generate a geometry file.""" + with open(output_path, "wb") as f: + f.write(self._bytes) + + +def _generate_meshes_from_polygons( + polygons: List[Tuple[Polygon, Dict[str, Any]]] +) -> List[trimesh.Trimesh]: + """Creates a mesh out of a list of polygons.""" + meshes = [] + + # Trimesh's API require a list of vertices and a list of faces, where each + # face contains three indexes into the vertices list. Ideally, the vertices + # are all unique and the faces list references the same indexes as needed. + # TODO: Batch the polygon processing. + for poly, metadata in polygons: + vertices, faces = [], [] + point_dict = dict() + current_point_index = 0 + + # Collect all the points on the shape to reduce checks by 3 times + for x, y in poly.exterior.coords: + p = (x, y, 0) + if p not in point_dict: + vertices.append(p) + point_dict[p] = current_point_index + current_point_index += 1 + triangles = triangulate_polygon(poly) + for triangle in triangles: + face = np.array( + [point_dict.get((x, y, 0), -1) for x, y in triangle.exterior.coords] + ) + # Add face if not invalid + if -1 not in face: + faces.append(face) + + if not vertices or not faces: + continue + + mesh = trimesh.Trimesh(vertices=vertices, faces=faces, metadata=metadata) + + # Trimesh doesn't support a coordinate-system="z-up" configuration, so we + # have to apply the transformation manually. + mesh.apply_transform( + trimesh.transformations.rotation_matrix(math.pi / 2, [-1, 0, 0]) + ) + meshes.append(mesh) + return meshes + + +def make_map_glb( + polygons: List[Tuple[Polygon, Dict[str, Any]]], + bbox: BoundingBox, + lane_dividers, + edge_dividers, +) -> GLBData: + """Create a GLB file from a list of road polygons.""" + scene = trimesh.Scene() + + # Attach additional information for rendering as metadata in the map glb + metadata = { + "bounding_box": ( + bbox.min_pt.x, + bbox.min_pt.y, + bbox.max_pt.x, + bbox.max_pt.y, + ), + "lane_dividers": lane_dividers, + "edge_dividers": edge_dividers, + } + + meshes = _generate_meshes_from_polygons(polygons) + for mesh in meshes: + mesh.visual = trimesh.visual.TextureVisuals( + material=trimesh.visual.material.PBRMaterial() + ) + + road_id = mesh.metadata["road_id"] + lane_id = mesh.metadata.get("lane_id") + name = str(road_id) + if lane_id is not None: + name += f"-{lane_id}" + scene.add_geometry(mesh, name, extras=mesh.metadata) + return GLBData(gltf.export_glb(scene, extras=metadata, include_normals=True)) + + +def make_road_line_glb(lines: List[List[Tuple[float, float]]]) -> GLBData: + """Create a GLB file from a list of road/lane lines.""" + scene = trimesh.Scene() + for line_pts in lines: + vertices = [(*pt, 0.1) for pt in line_pts] + point_cloud = trimesh.PointCloud(vertices=vertices) + point_cloud.apply_transform( + trimesh.transformations.rotation_matrix(math.pi / 2, [-1, 0, 0]) + ) + scene.add_geometry(point_cloud) + return GLBData(gltf.export_glb(scene)) diff --git a/smarts/core/waymo_map.py b/smarts/core/waymo_map.py index 568f3c4737..87bd98d19b 100644 --- a/smarts/core/waymo_map.py +++ b/smarts/core/waymo_map.py @@ -32,11 +32,10 @@ import numpy as np import rtree -import trimesh from cached_property import cached_property from shapely.geometry import Point as SPoint from shapely.geometry import Polygon -from trimesh.exchange import gltf +from smarts.core.utils.glb import make_map_glb, make_road_line_glb from waymo_open_dataset.protos import scenario_pb2 from waymo_open_dataset.protos.map_pb2 import ( Crosswalk, @@ -54,7 +53,7 @@ from .road_map import RoadMap, RoadMapWithCaches, Waypoint from .route_cache import RouteWithCache from .utils.file import read_tfrecord_file -from .utils.geometry import buffered_shape, generate_meshes_from_polygons +from .utils.geometry import buffered_shape from .utils.math import ( inplace_unwrap, line_intersect_vectorized, @@ -64,35 +63,6 @@ ) -def _convert_camera(camera): - result = { - "name": camera.name, - "type": "perspective", - "perspective": { - "aspectRatio": camera.fov[0] / camera.fov[1], - "yfov": np.radians(camera.fov[1]), - "znear": float(camera.z_near), - # HACK: The trimesh gltf export doesn't include a zfar which Panda3D GLB - # loader expects. Here we override to make loading possible. - "zfar": float(camera.z_near + 100), - }, - } - return result - - -gltf._convert_camera = _convert_camera - - -class _GLBData: - def __init__(self, bytes_): - self._bytes = bytes_ - - def write_glb(self, output_path: str): - """Generate a geometry file.""" - with open(output_path, "wb") as f: - f.write(self._bytes) - - class WaymoMap(RoadMapWithCaches): """A map associated with a Waymo dataset""" @@ -295,7 +265,7 @@ def _compute_lane_intersections(self, composites: bool): if lane._feature_id == cand_lane._feature_id: continue # Don't check intersection with incoming/outgoing lanes - if cand_id in lane.incoming_lanes or cand_id in lane.outgoing_lanes: + if cand_lane in lane.incoming_lanes or cand_lane in lane.outgoing_lanes: continue # ... or lanes in same road (TAI?) if lane.road == cand_lane.road: @@ -992,14 +962,8 @@ def scale_factor(self) -> float: def to_glb(self, glb_dir): """Build a glb file for camera rendering and envision.""" - glb = self._make_glb_from_polys() - glb.write_glb(Path(glb_dir) / "map.glb") - - def _make_glb_from_polys(self): - scene = trimesh.Scene() polygons = [] - for lane_id in self._lanes: - lane = self._lanes[lane_id] + for lane_id, lane in self._lanes.items(): metadata = { "road_id": lane.road.road_id, "lane_id": lane_id, @@ -1007,36 +971,13 @@ def _make_glb_from_polys(self): } polygons.append((lane.shape(), metadata)) - meshes = generate_meshes_from_polygons(polygons) - - # Attach additional information for rendering as metadata in the map glb - # <2D-BOUNDING_BOX>: four floats separated by ',' (,,,), - # which describe x-minimum, y-minimum, x-maximum, and y-maximum - metadata = { - "bounding_box": ( - self.bounding_box.min_pt.x, - self.bounding_box.min_pt.y, - self.bounding_box.max_pt.x, - self.bounding_box.max_pt.y, - ) - } - - # lane markings information lane_dividers = self._compute_traffic_dividers() - metadata["lane_dividers"] = lane_dividers - for mesh in meshes: - mesh.visual = trimesh.visual.TextureVisuals( - material=trimesh.visual.material.PBRMaterial() - ) + map_glb = make_map_glb(polygons, self.bounding_box, lane_dividers, []) + map_glb.write_glb(Path(glb_dir) / "map.glb") - road_id = mesh.metadata["road_id"] - lane_id = mesh.metadata.get("lane_id") - name = f"{road_id}" - if lane_id is not None: - name += f"-{lane_id}" - scene.add_geometry(mesh, name, extras=mesh.metadata) - return _GLBData(gltf.export_glb(scene, extras=metadata, include_normals=True)) + lane_lines_glb = make_road_line_glb(lane_dividers) + lane_lines_glb.write_glb(Path(glb_dir) / "lane_lines.glb") def _compute_traffic_dividers(self): lane_dividers = [] # divider between lanes with same traffic direction @@ -1933,7 +1874,8 @@ def _equally_spaced_path( lp_spacing: float, ) -> List[Waypoint]: """given a list of LanePoints starting near point, return corresponding - Waypoints that may not be evenly spaced (due to lane change) but start at point.""" + Waypoints that may not be evenly spaced (due to lane change) but start at point. + """ continuous_variables = [ "positions_x", @@ -1949,7 +1891,6 @@ def _equally_spaced_path( parameter: [] for parameter in (continuous_variables + discrete_variables) } for idx, lanepoint in enumerate(path): - if lanepoint.is_inferred and 0 < idx < len(path) - 1: continue diff --git a/smarts/sstudio/genhistories.py b/smarts/sstudio/genhistories.py index 7706a758fb..e3aa064d58 100644 --- a/smarts/sstudio/genhistories.py +++ b/smarts/sstudio/genhistories.py @@ -24,6 +24,7 @@ import logging import math import os +from pathlib import Path import sqlite3 import sys from collections import deque @@ -43,16 +44,6 @@ from smarts.sstudio import types from smarts.waymo.waymo_utils import WaymoDatasetError -try: - # pytype: disable=import-error - from waymo_open_dataset.protos import scenario_pb2 - from waymo_open_dataset.protos.map_pb2 import TrafficSignalLaneState - - # pytype: enable=import-error -except ImportError: - print( - "You may not have installed the [waymo] dependencies required to use the waymo replay simulation. Install them first using the command `pip install -e .[waymo]` at the source directory." - ) METERS_PER_FOOT = 0.3048 DEFAULT_LANE_WIDTH = 3.7 # a typical US highway lane is 12ft ~= 3.7m wide @@ -800,6 +791,15 @@ def __init__(self, dataset_spec: Dict[str, Any], output: str): super().__init__(dataset_spec, output) def _get_scenario(self): + try: + from waymo_open_dataset.protos import ( + scenario_pb2, + ) # pytype: disable=import-error + except ImportError: + print( + "You may not have installed the [waymo] dependencies required to use the waymo replay simulation. Install them first using the command `pip install -e .[waymo]` at the source directory." + ) + if "scenario_id" not in self._dataset_spec: errmsg = "Dataset spec requires scenario_id to be set" self._log.error(errmsg) @@ -939,6 +939,15 @@ def lerp(a: float, b: float, t: float) -> float: yield rows[j] def _encode_tl_state(self, waymo_state) -> SignalLightState: + try: + from waymo_open_dataset.protos.map_pb2 import ( + TrafficSignalLaneState, + ) # pytype: disable=import-error + except ImportError: + print( + "You may not have installed the [waymo] dependencies required to use the waymo replay simulation. Install them first using the command `pip install -e .[waymo]` at the source directory." + ) + if waymo_state == TrafficSignalLaneState.LANE_STATE_STOP: return SignalLightState.STOP if waymo_state == TrafficSignalLaneState.LANE_STATE_CAUTION: @@ -989,6 +998,103 @@ def column_val_in_row(self, row, col_name: str) -> Any: return row[col_name] +class Argoverse(_TrajectoryDataset): + """A tool for conversion of an Argoverse 2 dataset for use within SMARTS.""" + + def __init__(self, dataset_spec: Dict[str, Any], output: str): + super().__init__(dataset_spec, output) + + @property + def rows(self) -> Generator[Dict, None, None]: + try: + # pytype: disable=import-error + from av2.datasets.motion_forecasting.scenario_serialization import ( + load_argoverse_scenario_parquet, + ) + from av2.datasets.motion_forecasting.data_schema import ( + ObjectType as AvObjectType, + ) + + # pytype: enable=import-error + except ImportError: + print( + "You may not have installed the [argoverse] dependencies required to use the Argoverse 2 replay simulation. Install them first using the command `pip install -e .[argoverse]` at the source directory." + ) + + ALLOWED_TYPES = frozenset( + { + AvObjectType.VEHICLE, + AvObjectType.PEDESTRIAN, + AvObjectType.MOTORCYCLIST, + AvObjectType.CYCLIST, + AvObjectType.BUS, + } + ) + + def _lookup_agent_type(agent_type: AvObjectType) -> int: + # See decode_vehicle_type in traffic_history.py + if agent_type == AvObjectType.MOTORCYCLIST: + return 1 # motorcycle + elif agent_type == AvObjectType.VEHICLE: + return 2 # passenger + elif agent_type == AvObjectType.BUS: + return 3 # truck + elif agent_type in {AvObjectType.PEDESTRIAN, AvObjectType.CYCLIST}: + return 4 # pedestrian/bicycle + else: + return 0 # other + + input_dir = Path(self._dataset_spec["input_path"]) + scenario_id = input_dir.stem + parquet_file = input_dir / f"scenario_{scenario_id}.parquet" + scenario = load_argoverse_scenario_parquet(parquet_file) + + # Normalize to start at 0, and convert to milliseconds + timestamps = (scenario.timestamps_ns - scenario.timestamps_ns[0]) * 1e-6 + + # The ego vehicle has a string ID, so we need to give it a unique int ID + all_ids = [int(t.track_id) for t in scenario.tracks if t.track_id != "AV"] + ego_id = max(all_ids) + 1 + + for track in scenario.tracks: + # Only use dynamic objects + if track.object_type not in ALLOWED_TYPES: + continue + + if track.track_id == "AV": + is_ego = 1 + vehicle_id = ego_id + else: + is_ego = 0 + vehicle_id = int(track.track_id) + vehicle_type = _lookup_agent_type(track.object_type) + + for obj_state in track.object_states: + row = dict() + row["vehicle_id"] = vehicle_id + row["type"] = vehicle_type + row["sim_time"] = timestamps[obj_state.timestep] + row["position_x"] = obj_state.position[0] + row["position_y"] = obj_state.position[1] + row["heading_rad"] = constrain_angle( + (obj_state.heading - math.pi / 2) % (2 * math.pi) + ) + row["speed"] = np.linalg.norm(np.array(obj_state.velocity)) + row["lane_id"] = 0 + row["is_ego_vehicle"] = is_ego + + # Dimensions are not present in the Argoverse data. Setting these to 0 + # means default values for each vehicle type will be used. + # See TrafficHistory.decode_vehicle_type(). + row["length"] = 0 + row["height"] = 0 + row["width"] = 0 + yield row + + def column_val_in_row(self, row, col_name: str) -> Any: + return row[col_name] + + def import_dataset( dataset_spec: types.TrafficHistoryDataset, output_path: str, @@ -1012,6 +1118,8 @@ def import_dataset( dataset = Interaction(dataset_dict, output) elif source == "Waymo": dataset = Waymo(dataset_dict, output) + elif source == "Argoverse": + dataset = Argoverse(dataset_dict, output) else: raise ValueError( f"unsupported TrafficHistoryDataset type: {dataset_spec.source_type}"