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

Resolve "Identify Map Segments (Intersection, Roundabout, ...) in Lanes (e.g. wod-motion)"

Closes #3

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

Author: sventarika

.gitignore

@@ -9,4 +9,6 @@ omegadoc
 *.html
 *.parquet
 *.xodr
-*.png
+*.png
+01_tracks.mcap
+testfile_wod.mcap

hdds-test.mcap

@@ -0,0 +1,53 @@
+import numpy as np
+import omega_prime
+from pathlib import Path
+from omega_prime.intersection_helper import *
+import sys
+
+# # Read all mcap files in the directory:
+# folder_path = Path("/scenario-center-playground/scenarios/")
+
+# for file in folder_path.glob("*.mcap"):
+#     print(f"Processing file: {file}")
+#     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)
+#     locator = omega_prime.Locator.from_maposi(r.map)
+#     mapsegment = omega_prime.mapsemgents.Mapsegmentation(r.map, locator)
+#     mapsegment.init_intersections()
+#     mapsegment.plot(output_path)
+#     mapsegment.plot_intersections(output_path)
+print(sys.getrecursionlimit())
+
+file = Path(
+    "/scenario-center-playground/data/training_20s/training_20s.tfrecord-00003-of-01000/training_20s.tfrecord-00003-of-01000_4a63430816f60102.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, compute_polygons=True, split_lanes=True, split_lanes_lenght=10)
+r.create_mapsegments()
+mapsegment = r.mapsegment
+locator = omega_prime.Locator.from_map(r.map)
+
+
+id = r.host_vehicle_idx
+# Assuming r.moving_objects is a list of objects with x and y attributes
+x_values = r.moving_objects[id].x.to_numpy()  # Convert Polars Series to NumPy array
+y_values = r.moving_objects[id].y.to_numpy()  # Convert Polars Series to NumPy array
+frame = r.moving_objects[id]._df["frame"].to_numpy()  # Convert Polars Series to NumPy array
+
+# Combine x and y into a single NumPy array of shape (n, 2)
+positions = np.column_stack((frame, x_values, y_values))
+
+xy = np.column_stack((x_values, y_values))
+sts = locator.xys2sts(xy)
+
+segments = mapsegment.trajectory_segment_detection(positions)
+mapsegment.plot(output_plot=output_path, trajectory=positions)
+mapsegment.plot_intersections(output_path)
+
+for i, segment in enumerate(segments):
+    if segment[1] is not None:
+        print(f"Segment {i}: {segment[1].idx} with the type {segment[1].type}")
+    else:
+        print(f"Segment {i}: Trajectory could not be assigned to a segment.")

intersection.py

@@ -6,10 +6,14 @@
 from .locator import LaneRelation, Locator
 from .map import Lane, LaneBoundary, Map, MapOsi
 from .recording import MovingObject, Recording
+from .mapsemgents import Mapsegmentation
+
+
 from importlib.metadata import version
 
 __version__ = version("omega_prime")
 
+
 __all__ = [
     "Recording",
     "MovingObject",
@@ -22,4 +26,5 @@
     "LaneRelation",
     "converters",
     "metrics",
+    "Mapsegmentation",
 ]

omega_prime/__init__.py

@@ -0,0 +1,81 @@
+import networkx as nx
+import numpy as np
+import matplotlib.pyplot as plt
+import shapely
+import shapely.ops
+
+
+def find_misclassified_intersection_lanes(lane_dict, all_lanes):
+    """
+    Finds lanes classified as `on_intersection` but have no intersections in the lane_dict.
+
+    Args:
+        lane_dict (dict): A dictionary where keys are lane indices and values are arrays of intersecting lane indices.
+        all_lanes (list): A list of all lane objects, each with an `idx` and `on_intersection` attribute.
+
+    Returns:
+        np.ndarray: An array of lane indices that are misclassified as `on_intersection`.
+    """
+    misclassified_lanes = []
+    for lane in all_lanes:
+        if lane.on_intersection and (lane.idx not in lane_dict or len(lane_dict[lane.idx]) == 0):
+            misclassified_lanes.append(lane.idx)
+    return np.array(misclassified_lanes)
+
+
+def add_lanexy_to_graph(G, lanes):
+    """
+    Adds lane coordinates to the graph as node attributes.
+
+    Args:
+        G (networkx.Graph): The graph to which lane coordinates will be added.
+        locator (omega_prime.Locator): The locator object containing lane information.
+
+    Returns:
+        networkx.Graph: The updated graph with lane coordinates as node attributes.
+    """
+    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
+    return G
+
+
+def plot_graph(G, Path):
+    """
+    Plots the graph with lane coordinates.
+
+    Args:
+        G (networkx.Graph): The graph to be plotted.
+        Path (str or Path): The file path to save the plot.
+    """
+    pos = {node: (data["x"], data["y"]) for node, data in G.nodes(data=True)}
+    nx.draw(G, pos, with_labels=True, node_size=10, font_size=5)
+    plt.title("Intersection Graph")
+    plt.xlabel("X Coordinate")  # Add label for the x-axis
+    plt.ylabel("Y Coordinate")  # Add label for the y-axis
+    plt.grid(True)  # Optional: Add a grid for better visualization
+    plt.savefig(Path)
+    plt.close()  # Close the plot to free memory
+
+
+def is_matrix_symmetrical(successors_dict, predecessors_dict):
+    # Check successors against predecessors
+    for lane_idx, successors in successors_dict.items():
+        for successor in successors:
+            if lane_idx not in predecessors_dict.get(successor, []):
+                print(f"Lane {lane_idx} is not in the predecessors of its successor {successor}")
+                print(f"Successors: {successors_dict[lane_idx]}")
+                print(f"Predecessors: {predecessors_dict.get(successor, [])}")
+                return False
+
+    # Check predecessors against successors
+    for lane_idx, predecessors in predecessors_dict.items():
+        for predecessor in predecessors:
+            if lane_idx not in successors_dict.get(predecessor, []):
+                print(f"Lane {lane_idx} is not in the successors of its predecessor {predecessor}")
+                print(f"Predecessors: {predecessors_dict[lane_idx]}")
+                print(f"Successors: {successors_dict.get(predecessor, [])}")
+                return False
+
+    return True

omega_prime/intersection_helper.py

@@ -38,6 +38,8 @@ def extend_linestring(cls, l: shapely.LineString, simplify: bool = True, l_appen
         if isinstance(l, shapely.Point):
             cl = np.asarray(l.coords)[0]
             return shapely.LineString(np.stack([cl - np.array([0, 1]), cl, cl + np.array([0, 1])]))
+        if l.length == 0:
+            return l
 
         startvec = np.asarray(l.interpolate(0, normalized=True).coords) - np.asarray(
             l.interpolate(0 + cls.epsi, normalized=True).coords
@@ -179,7 +181,7 @@ class LaneRelation(StrEnum):
     neighbour_left = "left neighbour"
 
 
-@dataclass
+@dataclass(repr=False)
 class Locator:
     all_lanes: Any  # array of all lanes
     external2internal_laneid: dict[Any, int] = field(init=False)
@@ -196,6 +198,7 @@ def from_map(cls, map):
         return cls(all_lanes=all_lanes)
 
     def __post_init__(self):
+        # Create mapping with lane_id as key
         self.external2internal_laneid = {l.idx: i for i, l in enumerate(self.all_lanes)}
         self.internal2external_laneid = [l.idx for l in self.all_lanes]
 
@@ -290,7 +293,7 @@ def _xys2sts(self, xys, polygons=None):
                     # Returns the indxes of all centerlines that are in range
                     nearby_idx = self.query_centerlines(poly, range_percentage=0.1)
                     # Connect the no_assosciation_idxs with the intersection_lane_ids
-                    no_asscociation_idxs = np.append(no_asscociation_idxs, [no_associations[idx]] * len(nearby_idx))
+                    no_asscociation_idxs = np.append(no_asscociation_idxs, [idx] * len(nearby_idx))
                     intersection_lane_ids = np.append(intersection_lane_ids, nearby_idx)
 
         # Need a convertion from float values to int values. This is because the shapely STRtree query_nearest returns float values
@@ -356,6 +359,7 @@ def _get_routing_graph(self):
         external2internal_laneid = self.external2internal_laneid
         g = nx.DiGraph()
         for lid, lane in enumerate(all_lanes):
+            g.add_node(lid, lane=lane)
             for external_pid in lane.predecessor_ids:
                 try:
                     g.add_edge(lid, external2internal_laneid[external_pid], label=LaneRelation.predecessor)
@@ -366,6 +370,8 @@ def _get_routing_graph(self):
                     g.add_edge(lid, external2internal_laneid[external_sid], label=LaneRelation.successor)
                 except KeyError:
                     pass
+            if lane.right_boundary is None or lane.left_boundary is None:
+                continue
             right_neigbours = [
                 int(i) for i in str_tree.query(lane.right_boundary.polyline, predicate="covered_by") if int(i) != lid
             ]
@@ -440,3 +446,7 @@ def get_single_lane_association(
 
     def __repr__(self):
         return f"Locator({len(self.all_lanes)} lanes)<{id(self)}>"
+
+    def update_lane_ids_dict(self):
+        self.external2internal_laneid = {l.idx: i for i, l in enumerate(self.all_lanes)}
+        self.internal2external_laneid = [l.idx for l in self.all_lanes]

omega_prime/locator.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass, field
 from typing import Any
-
+from collections import namedtuple
 import betterosi
 import numpy as np
 import shapely
@@ -11,6 +11,9 @@
 import polars_st as st
 
 
+OsiLaneId = namedtuple("OsiLaneId", ["road_id", "lane_id"])
+
+
 @dataclass
 class ProjectionOffset:
     x: float
@@ -133,6 +136,10 @@ def end_points(self):
 @dataclass(repr=False)
 class LaneOsiCenterline(LaneBase):
     _osi: betterosi.Lane
+    left_boundary = None
+    right_boundary = None
+    on_intersection: bool = field(init=False, default=None)
+    is_approaching: bool = field(init=False, default=None)
 
     @staticmethod
     def _get_centerline(lane: betterosi.Lane):
@@ -143,14 +150,21 @@ def _get_centerline(lane: betterosi.Lane):
 
     @classmethod
     def create(cls, lane: betterosi.Lane):
+        successor_ids = [
+            p.successor_lane_id.value for p in lane.classification.lane_pairing if p.successor_lane_id is not None
+        ]
+        predecessor_ids = [
+            p.antecessor_lane_id.value for p in lane.classification.lane_pairing if p.antecessor_lane_id is not None
+        ]
+        lid = lane.id.value
         return cls(
             _osi=lane,
-            idx=lane.id.value,
+            idx=OsiLaneId(road_id=lid, lane_id=lid),
             centerline=cls._get_centerline(lane),
             type=betterosi.LaneClassificationType(lane.classification.type),
             subtype=betterosi.LaneClassificationSubtype(lane.classification.subtype),
-            successor_ids=[],
-            predecessor_ids=[],
+            successor_ids=np.array(list(set(successor_ids))),
+            predecessor_ids=np.array(list(set(predecessor_ids))),
         )
 
 
@@ -162,9 +176,10 @@ class LaneOsi(Lane, LaneOsiCenterline):
 
     @classmethod
     def create(cls, lane: betterosi.Lane):
+        lid = int(lane.id.value)
         return cls(
             _osi=lane,
-            idx=int(lane.id.value),
+            idx=OsiLaneId(road_id=lid, lane_id=lid),
             centerline=cls._get_centerline(lane),
             type=betterosi.LaneClassificationType(lane.classification.type),
             subtype=betterosi.LaneClassificationSubtype(lane.classification.subtype),
@@ -296,7 +311,10 @@ def create(cls, gt: betterosi.GroundTruth):
         return cls(
             _osi=gt,
             lane_boundaries={b.id.value: LaneBoundaryOsi.create(b) for b in gt.lane_boundary},
-            lanes={l.id.value: LaneOsi.create(l) for l in gt.lane if len(l.classification.right_lane_boundary_id) > 0},
+            lanes={
+                l.idx: l
+                for l in [LaneOsi.create(l) for l in gt.lane if len(l.classification.right_lane_boundary_id) > 0]
+            },
         )
 
     def __post_init__(self):
@@ -305,7 +323,10 @@ def __post_init__(self):
     def setup_lanes_and_boundaries(self):
         for b in self.lane_boundaries.values():
             b._map = self
+        map_osi_id2idx = {l._osi.id.value: l.idx for l in self.lanes.values()}
         for l in self.lanes.values():
+            l.successor_ids = [map_osi_id2idx[i] for i in l.successor_ids]
+            l.predecessor_ids = [map_osi_id2idx[i] for i in l.predecessor_ids]
             l._map = self
             l.set_boundaries()
             l.set_polygon()
@@ -322,10 +343,28 @@ def create(cls, gt: betterosi.GroundTruth):
             raise RuntimeError("No Map")
         return cls(
             _osi=gt,
-            lanes={l.id.value: LaneOsiCenterline.create(l) for l in gt.lane},
+            lanes={l.idx: l for l in [LaneOsiCenterline.create(l) for l in gt.lane]},
             lane_boundaries={},
         )
 
     def setup_lanes_and_boundaries(self):
+        map_osi_id2idx = {l._osi.id.value: l.idx for l in self.lanes.values()}
+        for l in self.lanes.values():
+            l.successor_ids = [map_osi_id2idx[int(i)] for i in l.successor_ids if int(i) in map_osi_id2idx]
+            l.predecessor_ids = [map_osi_id2idx[int(i)] for i in l.predecessor_ids if int(i) in map_osi_id2idx]
         for l in self.lanes.values():
             l._map = self
+
+        # Sometimes a presuccessor lane is not set as a successor lane in the other lane, therefore we need to check where this is the case and add it
+        for l in self.lanes.values():
+            for p in l.predecessor_ids:
+                if p not in self.lanes:
+                    continue
+                if l.idx not in self.lanes[p].successor_ids:
+                    self.lanes[p].successor_ids.append(l.idx)
+        for l in self.lanes.values():
+            for s in l.successor_ids:
+                if s not in self.lanes:
+                    continue
+                if l.idx not in self.lanes[s].predecessor_ids:
+                    self.lanes[s].predecessor_ids.append(l.idx)