Fix argoverse map (#71)

* one solution

* seems greate

* fix bug

* calculate wieth

* fix lane width and format

metadrive/component/blocks/argoverse_block.py

@@ -1,4 +1,6 @@
-from typing import Iterable
+import copy
+from metadrive.component.lane.abs_lane import LineColor, LineType
+from typing import Dict
 
 from metadrive.component.blocks.base_block import BaseBlock
 from metadrive.component.lane.argoverse_lane import ArgoverseLane
@@ -7,14 +9,138 @@
 
 
 class ArgoverseBlock(BaseBlock):
-    def __init__(self, block_index: int, global_network: RoadNetwork, argoverse_lanes: Iterable[ArgoverseLane]):
+    def __init__(self, block_index: int, global_network: RoadNetwork, argoverse_lanes: Dict[int, ArgoverseLane]):
         """
         No randomization when create argoverse block, Split Argoverse Map to several blocks to boost efficiency
         """
         super(ArgoverseBlock, self).__init__(block_index, global_network, 0)
+        self.origin_block_network = None
         self.argo_lanes = argoverse_lanes
 
     def _sample_topology(self) -> bool:
-        for lane in self.argo_lanes:
-            self.block_network.add_road(Road(lane.start_node, lane.end_node), [lane])
+        self.add_lanes()
+        self.set_width_in_intersect()
+        # for _from, _to_dict in self.block_network.graph.items():
+        #     for _to, lanes in _to_dict.items():
+        #         self.propagate_width(current_lanes=lanes)
         return True
+
+    # def _create_in_world(self):
+    #     # redundant_road_to_remove = []
+    #     # self.origin_block_network = copy.copy(self.block_network)
+    #     # for _from, dest in self.block_network.graph.items():
+    #     #     for _to, lanes in dest.items():
+    #     #         for lane in lanes:
+    #     #             if lane.start_node != _from and lane.end_node != _to:
+    #     #                 redundant_road_to_remove.append(Road(lane.start_node, lane.end_node))
+    #     # for road in redundant_road_to_remove:
+    #     #     self.block_network.remove_road(road)
+    #     return super(ArgoverseBlock, self)._create_in_world()
+
+    def set_width_in_intersect(self):
+        for lane in self.argo_lanes.values():
+            if lane.is_intersection:
+                if lane.width != lane.LANE_WIDTH:
+                    continue
+                if lane.successors is not None:
+                    for next_lane_id in lane.successors:
+                        if next_lane_id in self.argo_lanes and self.argo_lanes[next_lane_id].width != lane.LANE_WIDTH:
+                            lane.width = self.argo_lanes[next_lane_id].width
+                            break
+                if lane.width != lane.LANE_WIDTH:
+                    continue
+                if lane.predecessors is not None:
+                    for next_lane_id in lane.predecessors:
+                        if next_lane_id in self.argo_lanes and self.argo_lanes[next_lane_id].width != lane.LANE_WIDTH:
+                            lane.width = self.argo_lanes[next_lane_id].width
+                            break
+
+    def add_lanes(self):
+        for lane in self.argo_lanes.values():
+            ret = [lane.id]
+            # find left neighbour
+            left_id = lane.l_neighbor_id if lane.l_neighbor_id in self.argo_lanes else None
+            in_same_dir = True if left_id is not None and lane.is_in_same_direction(self.argo_lanes[left_id]) else False
+            while left_id is not None and in_same_dir:
+                ret.insert(0, left_id)
+                left_id = self.argo_lanes[left_id].l_neighbor_id
+                in_same_dir = True if left_id is not None and lane.is_in_same_direction(
+                    self.argo_lanes[left_id]
+                ) else False
+
+            # find right neighbour
+            right_id = lane.r_neighbor_id if lane.r_neighbor_id in self.argo_lanes else None
+            in_same_dir = True if right_id is not None and lane.is_in_same_direction(
+                self.argo_lanes[right_id]
+            ) else False
+            while right_id is not None and in_same_dir:
+                ret.append(right_id)
+                right_id = self.argo_lanes[right_id].r_neighbor_id
+                in_same_dir = True if right_id is not None and lane.is_in_same_direction(
+                    self.argo_lanes[right_id]
+                ) else False
+            lanes = [self.argo_lanes[id] for id in ret]
+            for idx, l in enumerate(lanes):
+                if l.is_intersection:
+                    # if l.turn_direction == "RIGHT" and l.r_neighbor_id is None:
+                    #     l.line_types = [LineType.NONE, LineType.CONTINUOUS]
+                    # else:
+                    #     l.line_types = [LineType.NONE, LineType.NONE]
+                    #
+                    # if l.turn_direction == "LEFT" and l.l_neighbor_id is None:
+                    #     l.line_types = [LineType.CONTINUOUS, LineType.NONE]
+                    # else:
+                    l.line_types = [LineType.NONE, LineType.NONE]
+                else:
+                    if l.r_neighbor_id is not None:
+                        right_type = LineType.BROKEN
+                    else:
+                        right_type = LineType.CONTINUOUS
+                    if idx == 0:
+                        left_type = LineType.CONTINUOUS
+                        if l.l_neighbor_id is not None:
+                            l.line_color = [LineColor.YELLOW, LineColor.GREY]
+                    else:
+                        left_type = LineType.BROKEN
+                    l.line_types = [left_type, right_type]
+            # if not lane.is_intersection:
+            _, right_lat = self.argo_lanes[lane.r_neighbor_id].local_coordinates(
+                lane.center_line_points[0]
+            ) if lane.r_neighbor_id in self.argo_lanes else (0, 0)
+            if (lane.l_neighbor_id in self.argo_lanes
+                    and not lane.is_intersection) or (lane.l_neighbor_id in self.argo_lanes and lane.is_intersection and
+                                                      self.argo_lanes[lane.l_neighbor_id].is_in_same_direction(lane)):
+                _, left_lat = self.argo_lanes[lane.l_neighbor_id].local_coordinates(lane.center_line_points[0])
+            else:
+                left_lat = 0
+            width = max(abs(right_lat), abs(left_lat)) + 0.1
+            lane.width = width if abs(width) > 1.0 else lane.LANE_WIDTH  # else default
+            self.block_network.add_road(Road(lane.start_node, lane.end_node), lanes)
+
+    # def propagate_width(self, current_lanes):
+    #     for current_lane in current_lanes:
+    #         lane_ids = current_lane.successors or []
+    #         # lane_ids += lane.predecessors if lane.predecessors is None else []
+    #         while len(lane_ids) > 0:
+    #             id = lane_ids.pop()
+    #             if id not in self.argo_lanes:
+    #                 continue
+    #             lane = self.argo_lanes[id]
+    #             if lane.l_neighbor_id is not None or lane.r_neighbor_id is not None:
+    #                 continue
+    #             if lane.width == lane.LANE_WIDTH:
+    #                 lane.width = current_lane.width
+    #                 lane_ids += lane.successors if lane.successors is not None else []
+    #
+    #         lane_ids = current_lane.predecessors or []
+    #         # lane_ids += lane.predecessors if lane.predecessors is None else []
+    #         while len(lane_ids) > 0:
+    #             id = lane_ids.pop()
+    #             if id not in self.argo_lanes:
+    #                 continue
+    #             lane = self.argo_lanes[id]
+    #             if lane.l_neighbor_id is not None or lane.r_neighbor_id is not None:
+    #                 continue
+    #             if lane.width == lane.LANE_WIDTH:
+    #                 lane.width = current_lane.width
+    #                 lane_ids += lane.predecessors if lane.predecessors is not None else []

metadrive/component/blocks/base_block.py

@@ -273,7 +273,7 @@ def _add_lane(self, lane: AbstractLane, lane_id: int, colors: List[Vec4]):
             for c, i in enumerate([-1, 1]):
                 line_color = colors[c]
                 acc_length = 0
-                if lane.line_types[c] != LineType.NONE:
+                if lane.line_types[c] == LineType.CONTINUOUS:
                     for segment in lane.segment_property:
                         lane_start = lane.position(acc_length, i * lane_width / 2)
                         acc_length += segment["length"]

metadrive/component/lane/argoverse_lane.py

@@ -9,7 +9,7 @@
 
 class ArgoverseLane(WayPointLane, LaneSegment):
     # according to api of get_vector_map_lane_polygons(), the lane width in argoverse dataset is 3.8m
-    LANE_WIDTH = 3.0
+    LANE_WIDTH = 3.8
 
     def __init__(
         self,
@@ -24,24 +24,30 @@ def __init__(
         predecessors: List[int],
         successors: Optional[List[int]],
         centerline: np.ndarray,
+        lane_width=None
     ):
         # convert_to_MetaDrive_coordinates
         centerline[:, 1] *= -1
         LaneSegment.__init__(
             self, id, has_traffic_control, turn_direction, is_intersection, l_neighbor_id, r_neighbor_id, predecessors,
             successors, centerline
         )
-        WayPointLane.__init__(self, centerline, self.LANE_WIDTH)
+        WayPointLane.__init__(self, centerline, self.LANE_WIDTH if lane_width is None else lane_width)
         self.start_node = start_node
         self.end_node = end_node
-        if is_intersection:
-            # if turn_direction == "RIGHT" and r_neighbor_id is None:
-            #     self.line_types = (LineType.NONE, LineType.CONTINUOUS)
-            # elif turn_direction == "LEFT" and l_neighbor_id is None:
-            #     self.line_types = (LineType.CONTINUOUS,LineType.NONE)
-            # else:
-            #     self.line_types = (LineType.CONTINUOUS if l_neighbor_id is None else LineType.NONE, LineType.CONTINUOUS if r_neighbor_id is None else LineType.NONE)
-            self.line_types = (LineType.NONE, LineType.NONE)
-        else:
-            self.line_types = (LineType.CONTINUOUS, LineType.CONTINUOUS)
+        # if is_intersection:
+        #     # if turn_direction == "RIGHT" and r_neighbor_id is None:
+        #     #     self.line_types = (LineType.NONE, LineType.CONTINUOUS)
+        #     # elif turn_direction == "LEFT" and l_neighbor_id is None:
+        #     #     self.line_types = (LineType.CONTINUOUS,LineType.NONE)
+        #     # else:
+        #     #     self.line_types = (LineType.CONTINUOUS if l_neighbor_id is None else LineType.NONE, LineType.CONTINUOUS if r_neighbor_id is None else LineType.NONE)
+        #     self.line_types = (LineType.NONE, LineType.NONE)
+        # else:
+        self.line_types = (LineType.CONTINUOUS, LineType.CONTINUOUS)
         self.line_color = (LineColor.GREY, LineColor.GREY)
+
+    def is_in_same_direction(self, another_lane):
+        if self.id == another_lane.l_neighbor_id and self.l_neighbor_id == another_lane.id:
+            return False
+        return super(ArgoverseLane, self).is_in_same_direction(another_lane)

metadrive/component/lane/waypoint_lane.py

@@ -1,4 +1,5 @@
 import math
+from metadrive.utils.math_utils import wrap_to_pi
 from typing import Tuple, Union
 
 import numpy as np
@@ -32,6 +33,8 @@ def __init__(
         # Segment is the part between two adjacent way points
         self.segment_property = self._get_properties()
         self.length = sum([seg["length"] for seg in self.segment_property])
+        self.is_straight = True if abs(self.heading_theta_at(0.1) -
+                                       self.heading_theta_at(self.length - 0.1)) < np.deg2rad(10) else False
 
     def _get_properties(self):
         ret = []
@@ -69,6 +72,9 @@ def width_at(self, longitudinal: float) -> float:
         return self.width
 
     def heading_theta_at(self, longitudinal: float) -> float:
+        """
+        In rad
+        """
         accumulate_len = 0
         for seg in self.segment_property:
             accumulate_len += seg["length"]
@@ -110,3 +116,17 @@ def segment(self, longitudinal: float):
             if accumulate_len + 0.1 >= longitudinal:
                 return self.segment_property[index]
         return self.segment_property[index]
+
+    def is_in_same_direction(self, another_lane):
+        """
+        Return True if two lane is in same direction
+        """
+        my_start_heading = self.heading_theta_at(0.1)
+        another_start_heading = another_lane.heading_theta_at(0.1)
+
+        my_end_heading = self.heading_theta_at(self.length - 0.1)
+        another_end_heading = another_lane.heading_theta_at(self.length - 0.1)
+
+        return True if abs(wrap_to_pi(my_end_heading) - wrap_to_pi(another_end_heading)) < 0.2 and abs(
+            wrap_to_pi(my_start_heading) - wrap_to_pi(another_start_heading)
+        ) < 0.2 else False

metadrive/component/map/argoverse_map.py

@@ -81,7 +81,7 @@ def _construct_road_network(self, lanes: list):
         for lane in lanes:
             self._post_process_lane(lane)
 
-        block = ArgoverseBlock(0, self.road_network, chosen_lanes)
+        block = ArgoverseBlock(0, self.road_network, {lane.id: lane for lane in chosen_lanes})
         block.construct_block(self.engine.worldNP, self.engine.physics_world)
         self.blocks.append(block)
 

metadrive/manager/real_data_manager.py

@@ -77,7 +77,7 @@ def _create_argoverse_vehicles_once(self, map: BaseMap) -> None:
         pos_dict = {i: j["init_pos"] for i, j in zip(locate_info.keys(), locate_info.values())}
 
         block = map.blocks[0]
-        lanes = block.argo_lanes
+        lanes = block.argo_lanes.values()
         roads = block.block_network.get_roads(direction='positive', lane_num=1)
         potential_vehicle_configs = []
         for l in lanes: