Merge branch 'Fix_lane_relations' into '3-identify-map-segments-intersection-roundabout-in-lanes-e-g-wod-motion'

Fix lane relations

See merge request fb-fi/data/omega-prime!11

Author: sventarika

intersection.py

@@ -19,10 +19,12 @@
 #     mapsegment.plot_intersections(output_path)
 print(sys.getrecursionlimit())
 
-file = Path("/scenario-center-playground/data/training_20s/training_20s_863518a3eb519031.mcap")
+file = Path(
+    "/scenario-center-playground/data/training_20s/training_20s.tfrecord-00106-of-01000/training_20s.tfrecord-00106-of-01000_8450f9adb20760a5.mcap"
+)
 output_path = Path("/scenario-center-playground/scenarios/") / file.stem
 output_path.mkdir(parents=True, exist_ok=True)
-r = omega_prime.Recording.from_file(filepath=file, split_lanes=True, split_lanes_lenght=5)
+r = omega_prime.Recording.from_file(filepath=file, compute_polygons=True, split_lanes=True, split_lanes_lenght=10)
 r.create_mapsegments()
 mapsegment = r.mapsegment
 locator = omega_prime.Locator.from_map(r.map)

omega_prime/intersection_helper.py

@@ -23,7 +23,7 @@ def find_misclassified_intersection_lanes(lane_dict, all_lanes):
     return np.array(misclassified_lanes)
 
 
-def add_lanexy_to_graph(G, locator):
+def add_lanexy_to_graph(G, lanes):
     """
     Adds lane coordinates to the graph as node attributes.
 
@@ -34,7 +34,7 @@ def add_lanexy_to_graph(G, locator):
     Returns:
         networkx.Graph: The updated graph with lane coordinates as node attributes.
     """
-    for lane in locator.all_lanes:
+    for lane in lanes.values():
         if lane.idx.lane_id in G.nodes:
             G.nodes[lane.idx.lane_id]["x"] = shapely.centroid(lane.centerline).x
             G.nodes[lane.idx.lane_id]["y"] = shapely.centroid(lane.centerline).y

omega_prime/mapsemgents.py

@@ -8,6 +8,7 @@
 from shapely.strtree import STRtree
 from shapely.geometry import Point
 from collections import namedtuple as nt
+from .locator import Locator
 
 concave_hull_ratio = 0.3
 
@@ -17,10 +18,10 @@ class Mapsegmentation:
     A class that handels multiple intersections on a single map.
     """
 
-    def __init__(self, map, locator, lane_buffer=None, intersection_overlap_buffer=None):
-        self.map = map
-        self.locator = locator
-        self.lanes = locator.all_lanes
+    def __init__(self, recording, lane_buffer=None, intersection_overlap_buffer=None):
+        self.map = recording.map
+        self.locator = Locator.from_map(recording.map)
+        self.lanes = recording.map.lanes
         self.intersections = []
         self.lane_dict = {}
         self.lane_successors_dict = {}
@@ -37,7 +38,7 @@ def __init__(self, map, locator, lane_buffer=None, intersection_overlap_buffer=N
         self.segments = []
 
         segment_name = nt("SegmentName", ["lane_id", "segment_idx", "segment"])
-        for lane in self.locator.all_lanes:
+        for lane in self.lanes.values():
             self.lane_segment_dict[lane.idx.lane_id] = segment_name(lane.idx.lane_id, None, None)
 
     def init_intersections(self):
@@ -53,7 +54,7 @@ def init_intersections(self):
         self.parallel_lane_dict = self.create_parallel_lane_dict()
         self.get_intersecting_lanes()
         self.graph_intersection_detection()
-        self.G = add_lanexy_to_graph(self.G, self.locator)
+        self.G = add_lanexy_to_graph(self.G, self.lanes)
         self.set_intersection_idx()
 
         if self.do_combine_intersections:
@@ -68,6 +69,7 @@ def init_intersections(self):
         self.set_lane_intersection_relation()
         self.check_if_all_lanes_are_on_segment()
         self.update_segment_ids()
+        self.create_lane_segment_dict()
         self.update_road_ids()
 
         # from pathlib import Path
@@ -78,12 +80,10 @@ def update_road_ids(self):
         """
         Updates the road_ids of the lane to the segment ID
         """
-        for i, lane in enumerate(self.locator.all_lanes):
+        for i, lane in enumerate(self.lanes.values()):
             lane_id = lane.idx.lane_id
             if lane_id in self.lane_segment_dict and self.lane_segment_dict[lane_id].segment is not None:
-                self.locator.all_lanes[i].idx = OsiLaneId(
-                    road_id=self.lane_segment_dict[lane_id].segment_idx, lane_id=lane_id
-                )
+                lane.idx = OsiLaneId(road_id=self.lane_segment_dict[lane_id].segment.idx, lane_id=lane_id)
             else:
                 # print(f"Lane {lane_id} is not part of an intersection, so it has no segment ID")
                 pass
@@ -100,18 +100,18 @@ def create_parallel_lane_dict(self):
         Returns:
             dict: A dictionary mapping each lane's lane_id to a list of parallel lane ids.
         """
-        lane_dict = {lane.idx.lane_id: [] for lane in self.locator.all_lanes}
+        lane_dict = {lane.idx.lane_id: [] for lane in self.lanes.values()}
 
         # Precompute lane directions for faster comparisons
         lane_directions = {}
-        for lane in self.locator.all_lanes:
+        for lane in self.lanes.values():
             coords = np.array(lane.centerline.coords)
             direction = coords[-1] - coords[0]
             lane_directions[lane.idx.lane_id] = direction / np.linalg.norm(direction)
 
         # Use spatial index to find potential parallel lanes
-        lane_centerlines = [lane.centerline for lane in self.locator.all_lanes]
-        lane_ids = [lane.idx.lane_id for lane in self.locator.all_lanes]
+        lane_centerlines = [lane.centerline for lane in self.lanes.values()]
+        lane_ids = [lane.idx.lane_id for lane in self.lanes.values()]
 
         # Only build tree if there are lanes
         if lane_centerlines:
@@ -120,7 +120,7 @@ def create_parallel_lane_dict(self):
             tree = STRtree(buffered_centerlines)
 
             # For each lane, find potential parallel lanes
-            for i, lane in enumerate(self.locator.all_lanes):
+            for i, lane in enumerate(self.lanes.values()):
                 lane_id = lane.idx.lane_id
                 buffer_geom = buffered_centerlines[i]
 
@@ -155,7 +155,7 @@ def check_if_all_lanes_are_on_segment(self):
         Returns:
             bool: True if all lanes are on a segment, False otherwise.
         """
-        for lane in self.locator.all_lanes:
+        for lane in self.lanes.values():
             lane_id = lane.idx.lane_id
             if lane_id not in self.lane_segment_dict or self.lane_segment_dict[lane_id].segment is None:
                 print(f"Lane {lane_id} is not on a segment")
@@ -192,9 +192,9 @@ def trajectory_segment_detection(self, trajectory):
                 buffered_polygons[segment.idx] = segment.polygon.buffer(buffer)
 
         for i, (frame, x, y, _, _) in enumerate(trajectory):
-            point = Point(x, y)
+            point = Point(x, y).buffer(2)
             for segment_id, buffered_poly in buffered_polygons.items():
-                if buffered_poly.contains(point):
+                if buffered_poly.intersects(point):
                     trajectory[i, 4] = segment_id
                     break
 
@@ -221,7 +221,7 @@ def create_lane_dict(self):
         Returns:
             lane_dict (dict): A dictionary mapping each lane's lane_id to the lane object.
         """
-        self.lane_dict = {lane.idx.lane_id: lane for lane in self.locator.all_lanes}
+        self.lane_dict = {lane.idx.lane_id: lane for lane in self.lanes.values()}
         return self.lane_dict
 
     def get_lane_successors_and_predecessors(self):
@@ -240,7 +240,7 @@ def get_lane_successors_and_predecessors(self):
         lane_successors = {}
         lane_predecessors = {}
 
-        for lane in self.locator.all_lanes:
+        for lane in self.lanes.values():
             # Use lane_id as keys, convert successor_ids to lane_ids if needed
             lane_successors[lane.idx.lane_id] = [
                 succ_id.lane_id if hasattr(succ_id, "lane_id") else succ_id for succ_id in lane.successor_ids
@@ -267,7 +267,7 @@ def get_intersecting_lanes(self, buffer: float = None):
             buffer = self.lane_buffer
 
         # Precompute buffered geometries
-        buffered_lanes = {lane.idx.lane_id: lane.centerline.buffer(buffer) for lane in self.lanes}
+        buffered_lanes = {lane.idx.lane_id: lane.centerline.buffer(buffer) for lane in self.lanes.values()}
 
         # Create spatial index
         buffered_geoms = list(buffered_lanes.values())
@@ -278,7 +278,7 @@ def get_intersecting_lanes(self, buffer: float = None):
             tree = None
 
         intersecting_lanes = {}
-        for i, lane in enumerate(self.lanes):
+        for i, lane in enumerate(self.lanes.values()):
             lane_id = lane.idx.lane_id
             if tree is None:
                 intersecting_lanes[lane_id] = []
@@ -454,7 +454,7 @@ def add_non_intersecting_lanes_to_intersection(self):
         """
         for intersection in self.intersections:
             intersection.update_polygon()
-            for lane in self.locator.all_lanes:
+            for lane in self.lanes.values():
                 lane_id = lane.idx.lane_id
                 if (
                     lane_id not in intersection.lane_ids
@@ -470,23 +470,25 @@ def create_lane_segment_dict(self):
         segment_name = nt("SegmentName", ["lane_id", "segment_idx", "segment"])
         # Combine self.intersections and self.isolated_connections into the lane_segment_dict
         segent_list = self.intersections + self.isolated_connections
+        lane_segment_dict = {lane_id: segment_name(lane_id, None, None) for lane_id in self.lane_dict.keys()}
 
         for segment in segent_list:
             for lane in segment.lanes:
                 lane_id = lane.idx.lane_id
                 try:
-                    if self.lane_segment_dict[lane_id].segment is None:
+                    if lane_segment_dict[lane_id].segment is None:
                         # The lane is not connected to any segment yet, so we can add it
-                        self.lane_segment_dict[lane_id] = segment_name(lane_id, segment.idx, segment)
+                        lane_segment_dict[lane_id] = segment_name(lane_id, segment.idx, segment)
                     else:
                         # If the lane is already in the dictionary, check if the segment is the same
-                        if self.lane_segment_dict[lane_id].segment_idx != segment.idx:
+                        if lane_segment_dict[lane_id].segment_idx != segment.idx:
                             raise ValueError(
-                                f"Lane {lane_id} is already in segment {self.lane_segment_dict[lane_id].segment_idx} but trying to add it to segment {segment.idx}"
+                                f"Lane {lane_id} is already in segment {lane_segment_dict[lane_id].segment_idx} but trying to add it to segment {segment.idx}"
                             )
                 except ValueError as e:
                     print(f"Error: {e}")
                     continue
+        self.lane_segment_dict = lane_segment_dict
 
     def create_non_intersecting_lane_graph(self):
         """Create a graph with each lane which is not part of a intersection as a node and the edges are the successors and predecessors of the lanes.
@@ -496,7 +498,7 @@ def create_non_intersecting_lane_graph(self):
             G (networkx.Graph): A graph with each lane as a node and the edges are the successors and predecessors of the lanes.
         """
         G = nx.Graph()
-        for lane in self.locator.all_lanes:
+        for lane in self.lanes.values():
             lane_id = lane.idx.lane_id
             if lane_id not in self.lane_segment_dict or self.lane_segment_dict[lane_id].segment is None:
                 G.add_node(lane_id)
@@ -572,10 +574,16 @@ def combine_isolated_connections(self, isolated_connections):
                 if (
                     i != j
                     and connection.intersection_idxs == connection2.intersection_idxs
-                    and len(connection.intersection_idxs) > 0
+                    and len(connection.intersection_idxs) > 1
+                ):
+                    combined_connections.append([i, j])
+                elif (
+                    i != j
+                    and connection.intersection_idxs == connection2.intersection_idxs
+                    and len(connection.intersection_idxs) == 1
+                    and connection.polygon.distance(connection2.polygon) < 5
                 ):
                     combined_connections.append([i, j])
-                    # print(f"Connection {i} and {j} can be combined")
         final_combined = self.find_resulting_intersections(combined_connections)
         # print(f"Final combined: {final_combined}")
         new_connections = []
@@ -606,21 +614,21 @@ def set_lane_intersection_relation(self):
         Sets the attribute lane.approching true if the lane is connecting to an intersection.
         Setss the attribute lane.on_intersection true if the lane is part of an intersection.
         """
-        for lane in self.locator.all_lanes:
-            lane.on_intersection = False
-            lane.is_approaching = False
+        for lane in self.lanes.values():
+            self.lanes[lane.idx].on_intersection = False
+            self.lanes[lane.idx].is_approaching = False
 
         for intersection in self.intersections:
             for lane in intersection.lanes:
-                lane.on_intersection = True
-                lane.is_approaching = False
+                self.lanes[lane.idx].on_intersection = True
+                self.lanes[lane.idx].is_approaching = False
 
             for predecessor_id in lane.predecessor_ids:
                 # Convert to lane_id if it's a namedtuple
                 pred_id = predecessor_id.lane_id if hasattr(predecessor_id, "lane_id") else predecessor_id
                 predecessor = self.lane_dict[pred_id]
                 if predecessor.on_intersection is not True:
-                    predecessor.is_approaching = True
+                    self.lanes[predecessor.idx].is_approaching = True
 
     def plot(self, output_plot=None, trajectory=None, plot_lane_ids=False):
         """
@@ -634,7 +642,7 @@ def plot(self, output_plot=None, trajectory=None, plot_lane_ids=False):
         fig, ax = plt.subplots(1, 1)
         ax.set_aspect(1)
 
-        for lane in self.locator.all_lanes:
+        for lane in self.lanes.values():
             c = "black" if not lane.type == betterosi.LaneClassificationType.TYPE_INTERSECTION else "green"
             ax.plot(*lane.centerline.xy, color=c, alpha=0.3, zorder=-10)
             if plot_lane_ids:

omega_prime/recording.py

@@ -23,7 +23,6 @@
 import altair as alt
 import polars as pl
 import polars_st as st
-from .locator import Locator
 from .mapsemgents import Mapsegmentation
 
 pi_valued = pa.Check.between(-np.pi, np.pi)
@@ -617,7 +616,7 @@ def from_file(
                                 lanes_new[new_lane.idx.lane_id] = new_lane
                                 idx_count += 1
 
-                                # Update references in other lanes' predecessors/successors
+                            # Update references in other lanes' predecessors/successors
                             for other_lane in lanes_or.values():
                                 if lane.idx in other_lane.successor_ids:
                                     # Replace reference to original lane with first segment
@@ -629,6 +628,46 @@ def from_file(
                                     other_lane.predecessor_ids[idx] = segment_lanes[-1].idx
 
                         # Replace original lanes with segmented lanes
+
+                        # Do a check for the presessor and successor: Check if the distance between the centerlines is greater than the max_len --> if yes, then remove the connection
+                        for lane in lanes_new.values():
+                            if lane.predecessor_ids:
+                                for pre in lane.predecessor_ids:
+                                    pre_to_remove = []
+                                    if lanes_new[pre.lane_id].centerline.distance(lane.centerline) > max_len:
+                                        print(
+                                            f"Removing connection between {pre} and {lane.idx} due to distance > {max_len}"
+                                        )
+                                        pre_to_remove.append(pre)
+                                        try:
+                                            lanes_new[pre.lane_id].successor_ids.remove(lane.idx)
+                                        except ValueError:
+                                            pass  # If the successor is not in the list, ignore
+
+                                    for pre in pre_to_remove:
+                                        try:
+                                            lanes_new[lane.idx.lane_id].predecessor_ids.remove(pre)
+                                        except ValueError:
+                                            pass  # If the predecessor is not in the list, ignore
+                            if lane.successor_ids:
+                                for suc in lane.successor_ids:
+                                    suc_to_remove = []
+                                    if lanes_new[suc.lane_id].centerline.distance(lane.centerline) > max_len:
+                                        print(
+                                            f"Removing connection between {lane.idx} and {suc} due to distance > {max_len}"
+                                        )
+                                        suc_to_remove.append(suc)
+                                        try:
+                                            lanes_new[suc.lane_id].predecessor_ids.remove(lane.idx)
+                                        except ValueError:
+                                            pass  # If the predecessor is not in the list, ignore
+
+                                    for suc in suc_to_remove:
+                                        try:
+                                            lanes_new[lane.idx.lane_id].successor_ids.remove(suc)
+                                        except ValueError:
+                                            pass
+
                         r.map.lanes = {lane.idx: lane for lane in lanes_new.values()}
                         for lane in r.map.lanes.values():
                             lane._map = r.map
@@ -895,6 +934,6 @@ def plot_altair(
         return view.add_params(op_var)
 
     def create_mapsegments(self):
-        locator = Locator.from_map(self.map)
-        self.mapsegment = Mapsegmentation(self.map, locator=locator)
+        self.mapsegment = Mapsegmentation(self)
         self.mapsegment.init_intersections()
+        print("test")