API

Recording

Class representing a continuous traffic observation. Usually corresponds to one omega-prime file.

Internally, the Recording uses a Polars DataFrame to store moving object data. Each row in the DataFrame represents the state of a moving object at a specific timestamp.

Attributes:

Name	Type	Description
df	DataFrame	Polars DataFrame containing the moving object data.
map	MapOsi | MapOsiCenterline | MapOdr | None	Map associated with the recording.
projections	dict	Projection metadata with structure {"proj_string": str | None, None: ProjectionOffset | None, int: ProjectionOffset | None}.
traffic_light_states	dict	Dictionary mapping timestamps to traffic light states.
host_vehicle_idx	int | None	Index of the host vehicle, if applicable.

Source code in omega_prime/recording.py

class Recording:
    """Class representing a continuous traffic observation. Usually corresponds to one omega-prime file.

    Internally, the Recording uses a Polars DataFrame to store moving object data. Each row in the DataFrame
    represents the state of a moving object at a specific timestamp.

    Attributes:
        df (pl.DataFrame): Polars DataFrame containing the moving object data.
        map (MapOsi | MapOsiCenterline | MapOdr | None): Map associated with the recording.
        projections (dict): Projection metadata with structure
            `{"proj_string": str | None, None: ProjectionOffset | None, int: ProjectionOffset | None}`.
        traffic_light_states (dict): Dictionary mapping timestamps to traffic light states.
        host_vehicle_idx (int | None): Index of the host vehicle, if applicable.
    """

    _MovingObjectClass: typing.ClassVar = MovingObject

    @staticmethod
    def _offset_components(
        offset: ProjectionOffset | None,
    ) -> tuple[float, float, float, float]:
        "Extract components from a ProjectionOffset, returning zeros if the offset is None."
        if offset is None:
            return 0.0, 0.0, 0.0, 0.0
        return offset.x, offset.y, offset.z, offset.yaw

    @staticmethod
    def _encode_projections(
        projections: dict[typing.Any, typing.Any],
    ) -> bytes:
        "Encode projection metadata into a JSON string and then to bytes"
        if not projections:
            return b""

        def _serialize_offset(offset: ProjectionOffset | None):
            if offset is None:
                return None
            return {"x": offset.x, "y": offset.y, "z": offset.z, "yaw": offset.yaw}

        payload = {
            "proj_string": projections.get("proj_string"),
            "offsets": [
                {
                    "total_nanos": ts,
                    "offset": _serialize_offset(offset),
                }
                for ts, offset in projections.items()
                if ts != "proj_string"
            ],
        }
        return json.dumps(payload).encode()

    @staticmethod
    def _decode_projections(
        raw: bytes | str | None,
    ) -> dict[typing.Any, typing.Any]:
        "Decode projection metadata from bytes or string to a dictionary."
        if raw in (None, b"", ""):
            return {}
        if isinstance(raw, bytes):
            raw = raw.decode()
        payload = json.loads(raw)
        result: dict[typing.Any, typing.Any] = {"proj_string": payload.get("proj_string")}
        for entry in payload.get("offsets", []):
            offset_data = entry.get("offset")
            offset = ProjectionOffset(**offset_data) if offset_data is not None else None
            ts = entry.get("total_nanos")
            key = None if ts is None else int(ts)
            result[key] = offset
        return result

    @staticmethod
    def _validate_projections_schema(
        projections: dict[typing.Any, typing.Any] | None,
    ) -> dict[typing.Any, typing.Any]:
        """
        Validate the schema of the projections dictionary, ensuring correct types and structure.
        Projection metadata with structure:
            `{"proj_string": str | None, None: ProjectionOffset | None, int: ProjectionOffset | None}`
        """
        if projections is None:
            return {}
        if not isinstance(projections, dict):
            raise TypeError("`projections` must be a dictionary.")

        validated: dict[typing.Any, typing.Any] = {}
        if "proj_string" in projections:
            validated["proj_string"] = projections["proj_string"]

        for key, value in projections.items():
            if key == "proj_string":
                continue

            if key is not None and not isinstance(key, int):
                raise TypeError("Projection keys must be integers, `None`, or `proj_string`.")

            if value is not None and not isinstance(value, ProjectionOffset):
                raise TypeError("Projection values must be `ProjectionOffset` or `None`.")

            validated[key] = value

        return validated

    def _projection_for_timestamp(self, total_nanos: int) -> tuple[str | None, ProjectionOffset | None]:
        source_proj_string = self.projections.get("proj_string")
        if source_proj_string is None:
            source_proj_string = getattr(self.map, "proj_string", None)

        if total_nanos in self.projections:
            offset = self.projections[total_nanos]
        elif None in self.projections:
            offset = self.projections[None]
        else:
            offset = None

        return source_proj_string, offset

    @staticmethod
    def get_moving_object_ground_truth(
        nanos: int,
        df: pl.DataFrame,
        host_vehicle_idx: int | None = None,
        validate: bool = False,
    ) -> betterosi.GroundTruth:
        if validate:
            recording_moving_object_schema.validate(df, lazy=True)

        def get_object(row):
            return betterosi.MovingObject(
                id=betterosi.Identifier(value=row["idx"]),
                type=betterosi.MovingObjectType(row["type"]),
                base=betterosi.BaseMoving(
                    dimension=betterosi.Dimension3D(length=row["length"], width=row["width"], height=row["width"]),
                    position=betterosi.Vector3D(x=row["x"], y=row["y"], z=row["z"]),
                    orientation=betterosi.Orientation3D(roll=row["roll"], pitch=row["pitch"], yaw=row["yaw"]),
                    velocity=betterosi.Vector3D(x=row["vel_x"], y=row["vel_y"], z=row["vel_z"]),
                    acceleration=betterosi.Vector3D(x=row["acc_x"], y=row["acc_y"], z=row["acc_z"]),
                ),
                vehicle_classification=betterosi.MovingObjectVehicleClassification(
                    type=row["subtype"], role=row["role"]
                ),
            )

        mvs = [get_object(r) for r in df.iter_rows(named=True)]
        gt = betterosi.GroundTruth(
            version=betterosi.InterfaceVersion(version_major=3, version_minor=7, version_patch=9),
            timestamp=betterosi.Timestamp(seconds=int(nanos // int(1e9)), nanos=int(nanos % int(1e9))),
            host_vehicle_id=(
                betterosi.Identifier(value=0)
                if host_vehicle_idx is None
                else betterosi.Identifier(value=host_vehicle_idx)
            ),
            moving_object=mvs,
        )
        return gt

    @staticmethod
    def _ensure_polars_dataframe(df: typing.Any) -> pl.DataFrame:
        "Ensure that the input data is a Polars DataFrame with the correct schema, converting if necessary."
        if isinstance(df, pl.DataFrame):
            return df
        return pl.DataFrame(df, schema_overrides=polars_schema)

    @staticmethod
    def _build_frame_mapping(df: pl.DataFrame) -> tuple[dict[int, int], pl.DataFrame]:
        "Build a mapping from `total_nanos` to frame numbers and return both the mapping and a DataFrame for joining."
        nanos2frame = {n: i for i, n in enumerate(df["total_nanos"].unique())}
        mapping = pl.DataFrame(
            {
                "total_nanos": list(nanos2frame.keys()),
                "frame": list(nanos2frame.values()),
            },
            schema=dict(total_nanos=polars_schema["total_nanos"], frame=pl.UInt32),
        )
        return nanos2frame, mapping

    @staticmethod
    def _attach_frame_column(df: pl.DataFrame, mapping: pl.DataFrame) -> pl.DataFrame:
        if "frame" in df.columns:
            df = df.drop("frame")
        return df.join(mapping, on="total_nanos", how="left")

    @staticmethod
    def _ensure_motion_norm_columns(df: pl.DataFrame) -> pl.DataFrame:
        exprs = []
        if "vel" not in df.columns:
            exprs.append((pl.col("vel_x") ** 2 + pl.col("vel_y") ** 2).sqrt().alias("vel"))
        if "acc" not in df.columns:
            exprs.append((pl.col("acc_x") ** 2 + pl.col("acc_y") ** 2).sqrt().alias("acc"))
        if exprs:
            df = df.with_columns(*exprs)
        return df

    def __init__(
        self,
        df,
        map=None,
        projections=None,
        host_vehicle_idx: int | None = None,
        validate=False,
        traffic_light_states: dict | None = None,
    ):
        "Initialize a Recording instance."
        df = self._ensure_polars_dataframe(df)
        if "total_nanos" not in df.columns:
            raise ValueError("df must contain column `total_nanos`.")
        nanos2frame, mapping = self._build_frame_mapping(df)
        df = self._attach_frame_column(df, mapping)
        df = self._ensure_polars_dataframe(df)
        if validate:
            recording_moving_object_schema.validate(df, lazy=True)

        super().__init__()
        self.nanos2frame = nanos2frame

        df = self._ensure_motion_norm_columns(df)
        self.projections = self._validate_projections_schema(projections)
        self.traffic_light_states = traffic_light_states if traffic_light_states is not None else {}

        df = bbx_to_polygon(df)

        self._df = df
        self.map = map
        self._moving_objects = None
        self.host_vehicle_idx = host_vehicle_idx
        self.mapsegment = None

    @property
    def df(self):
        return self._df

    @property
    def host_vehicle(self):
        return self.moving_objects.get(self.host_vehicle_idx, None)

    @property
    def moving_objects(self):
        if self._moving_objects is None:
            self._mv_df = (
                self._df.group_by("idx")
                .agg(
                    pl.col("length", "width", "height").mean(),
                    pl.col("type", "subtype", "role").median(),
                    pl.col("frame").min().alias("birth"),
                    pl.col("frame").max().alias("end"),
                    pl.col("total_nanos").min().alias("t_birth"),
                    pl.col("total_nanos").max().alias("t_end"),
                )
                .with_columns(
                    pl.col("type").map_elements(lambda x: betterosi.MovingObjectType(x), return_dtype=object),
                    pl.col("subtype").map_elements(
                        lambda x: (betterosi.MovingObjectVehicleClassificationType(x) if x != -1 else None),
                        return_dtype=object,
                    ),
                    pl.col("role").map_elements(
                        lambda x: (betterosi.MovingObjectVehicleClassificationRole(x).name if x != -1 else None),
                        return_dtype=object,
                    ),
                )
            )
            self._moving_objects = {int(idx): self._MovingObjectClass(self, idx) for idx in self._df["idx"].unique()}

        return self._moving_objects

    def _df_with_original_pose_for_export(self, df: pl.DataFrame | None = None) -> pl.DataFrame:
        """
        Return a DataFrame with original pose columns (`x_original`, `y_original`, `z_original`, `yaw_original`)
        for export, if they exist. This is used to ensure that the original pose information is preserved when
        exporting to formats like Parquet or MCAP,
        even if the main `x`, `y`, `z`, and `yaw` columns have been modified by projections.
        """
        df_export = self._df if df is None else df
        original_to_base = {
            "x_original": "x",
            "y_original": "y",
            "z_original": "z",
            "yaw_original": "yaw",
        }
        overwrite_exprs = [
            pl.col(original_col).alias(base_col)
            for original_col, base_col in original_to_base.items()
            if original_col in df_export.columns
        ]
        if overwrite_exprs:
            df_export = df_export.with_columns(*overwrite_exprs)
        return df_export

    def to_osi_gts(self) -> list[betterosi.GroundTruth]:
        first_iteration = True
        df_export = self._df_with_original_pose_for_export()
        for [nanos], group_df in df_export.sort(["total_nanos"]).group_by("total_nanos", maintain_order=True):
            gt = self.get_moving_object_ground_truth(
                nanos, group_df, host_vehicle_idx=self.host_vehicle_idx, validate=False
            )
            source_proj_string, proj_offset = self._projection_for_timestamp(int(nanos))
            if source_proj_string is not None:
                gt.proj_string = source_proj_string
            if proj_offset is not None:
                gt.proj_frame_offset = betterosi.GroundTruthProjFrameOffset(
                    position=betterosi.Vector3D(x=proj_offset.x, y=proj_offset.y, z=proj_offset.z),
                    yaw=proj_offset.yaw,
                )
            if first_iteration:
                first_iteration = False
                if self.map is not None and isinstance(self.map, MapOsi | MapOsiCenterline):
                    gt.lane_boundary = [b._osi for b in self.map.lane_boundaries.values()]
                    gt.lane = [l._osi for l in self.map.lanes.values()]
            if nanos in self.traffic_light_states:
                gt.traffic_light = self.traffic_light_states[nanos]
            yield gt

    @classmethod
    def from_osi_gts(cls, gts: list[betterosi.GroundTruth], **kwargs):
        projs: dict[typing.Any, typing.Any] = {"proj_string": None}
        traffic_light_states = {}

        gts, tmp_gts = itertools.tee(gts, 2)
        first_gt = next(tmp_gts)
        if first_gt.host_vehicle_id is not None:
            host_vehicle_idx = first_gt.host_vehicle_id.value
        else:
            host_vehicle_idx = None

        def get_gts():
            for i, gt in enumerate(gts):
                total_nanos = gt.timestamp.seconds * 1_000_000_000 + gt.timestamp.nanos
                if gt.proj_frame_offset is not None and gt.proj_frame_offset.position is None:
                    raise ValueError(
                        f"Offset of {i}th ground truth message (total_nanos={total_nanos}) is set without position."
                    )

                projs[total_nanos] = (
                    ProjectionOffset(
                        x=gt.proj_frame_offset.position.x,
                        y=gt.proj_frame_offset.position.y,
                        z=gt.proj_frame_offset.position.z,
                        yaw=gt.proj_frame_offset.yaw,
                    )
                    if gt.proj_frame_offset is not None
                    else None
                )

                if gt.proj_string is not None:
                    normalized_proj_string = gt.proj_string.strip()
                    if normalized_proj_string:
                        if projs["proj_string"] is None:
                            projs["proj_string"] = normalized_proj_string
                        elif projs["proj_string"] != normalized_proj_string:
                            raise ValueError(
                                f"Conflicting projection strings: {projs['proj_string']} vs {normalized_proj_string} at gt index {i} (total_nanos={total_nanos})."
                            )

                traffic_light_states[total_nanos] = gt.traffic_light

                for mv in gt.moving_object:
                    yield dict(
                        total_nanos=total_nanos,
                        idx=mv.id.value,
                        x=mv.base.position.x,
                        y=mv.base.position.y,
                        z=mv.base.position.z,
                        vel_x=mv.base.velocity.x,
                        vel_y=mv.base.velocity.y,
                        vel_z=mv.base.velocity.z,
                        acc_x=mv.base.acceleration.x,
                        acc_y=mv.base.acceleration.y,
                        acc_z=mv.base.acceleration.z,
                        length=mv.base.dimension.length,
                        width=mv.base.dimension.width,
                        height=mv.base.dimension.height,
                        roll=mv.base.orientation.roll,
                        pitch=mv.base.orientation.pitch,
                        yaw=mv.base.orientation.yaw,
                        type=mv.type,
                        role=(
                            mv.vehicle_classification.role if mv.type == betterosi.MovingObjectType.TYPE_VEHICLE else -1
                        ),
                        subtype=(
                            mv.vehicle_classification.type if mv.type == betterosi.MovingObjectType.TYPE_VEHICLE else -1
                        ),
                    )

        df_mv = pl.DataFrame(get_gts(), schema=polars_schema).sort(["total_nanos", "idx"])
        return cls(
            df_mv,
            projections=projs,
            host_vehicle_idx=host_vehicle_idx,
            traffic_light_states=traffic_light_states,
            **kwargs,
        )

    def to_mcap(self, filepath):
        "Store Recording as an MCAP file."
        if Path(filepath).suffix != ".mcap":
            raise ValueError()
        with betterosi.Writer(filepath) as w:
            for gt in self.to_osi_gts():
                w.add(gt)
            if isinstance(self.map, MapOdr):
                w.add(self.map.to_osi(), topic="ground_truth_map", log_time=0)
            elif (
                self.map is not None and not isinstance(self.map, MapOsi) and not isinstance(self.map, MapOsiCenterline)
            ):
                warn(f"The map {self.map} could not be saved to `mcap`")

    @classmethod
    def from_parquet(cls, filename, parse_map: bool = False, step_size: float = 0.01, **kwargs):
        t = pq.read_table(filename)
        df = pl.DataFrame(t, schema_overrides=polars_schema)
        host_vehicle_idx = None
        projections: dict[typing.Any, typing.Any] = {}
        map = None
        metadata = t.schema.metadata or {}
        if metadata:
            if b"host_vehicle_idx" in metadata:
                host_vehicle_idx = int(metadata[b"host_vehicle_idx"].decode())

            projections = cls._decode_projections(metadata.get(b"projections_json"))

            map_parsing = {}
            for MC in MAP_CLASSES:
                if MC._binary_json_identifier in metadata:
                    try:
                        map = MC._from_binary_json(
                            metadata,
                            parse_map=parse_map,
                            step_size=step_size,
                        )
                    except Exception as e:
                        map_parsing[MC.__name__] = str(e)
                    else:
                        if map is not None:
                            break

        return cls(
            df,
            map=map,
            host_vehicle_idx=host_vehicle_idx,
            projections=projections,
            **kwargs,
        )

    def to_parquet(self, filename):
        "Store Recording as a Parquet file."
        metadata = {}
        if self.host_vehicle_idx is not None:
            metadata[b"host_vehicle_idx"] = str(self.host_vehicle_idx).encode()
        proj_meta = {}
        encoded_projections = self._encode_projections(self.projections)
        if encoded_projections:
            proj_meta[b"projections_json"] = encoded_projections
        df_export = self._df_with_original_pose_for_export()
        to_drop = ["frame"]
        optional_cols = [
            "polygon",
            "global_lat",
            "global_lon",
            "global_alt",
            "global_yaw",
            "proj_string",
            "x_original",
            "y_original",
            "z_original",
            "yaw_original",
        ]
        to_drop.extend([c for c in optional_cols if c in df_export.columns])
        t = pyarrow.table(df_export.drop(*to_drop))
        map_meta = self.map._to_binary_json() if self.map is not None else {}

        t = t.cast(t.schema.with_metadata(metadata | proj_meta | map_meta))
        pq.write_table(t, filename)

    @classmethod
    def from_file(
        cls,
        filepath,
        map_path: str | None = None,
        validate: bool = False,
        parse_map: bool = False,
        step_size: float = 0.01,
        apply_proj: bool = True,
        **kwargs,
    ) -> "Recording":
        """Load a Recording from a file. Supports `.parquet`, `.osi` and `.mcap` files.

        Parameters:
            filepath (str): Path to the input file.
            map_path (str | None): Optional path to a map file. If None, the map will be loaded from the recording if available.
            validate (bool): Whether to validate the data against the schema.
            parse_map (bool): Whether to create python objects from the map data or just load it.
            step_size (float): Step size for map parsing, if applicable (Used for ASAM OpenDRIVE).
            apply_proj (bool): Whether to apply projection transformations to the recording's moving object data.

        Returns:
            Recording (Recording): The loaded Recording object.
        """
        if filepath is None and map_path is None:
            raise ValueError("Either `filepath` or `map_path` must be provided.")

        if filepath is not None and Path(filepath).suffix == ".parquet":
            r = cls.from_parquet(filepath, parse_map=parse_map, validate=validate, step_size=step_size)
        elif filepath is not None:
            gts = betterosi.read(filepath, return_ground_truth=True, mcap_return_betterosi=True)
            r = cls.from_osi_gts(gts, validate=validate)
        if map_path is None and r.map is not None:
            return r

        map_path = Path(map_path if map_path is not None else filepath)
        map_parsing = {}
        map = None
        for MC in MAP_CLASSES:
            if map_path.suffix in MC._supported_file_suffixes:
                try:
                    map = MC.from_file(map_path, parse_map=parse_map, **kwargs)
                except Exception as e:
                    map_parsing[MC.__name__] = str(e)
                else:
                    break
        if map is not None:
            r.map = map
        elif r.map is None:
            warn(f"No map could be found: {map_parsing}")

        if r.projections and apply_proj:
            try:
                r.apply_projections()
            except Exception:
                warn("Failed to apply projections.")
        return r

    def to_file(self, filepath):
        "Store Recording to a file based on its suffix (`.parquet`, `.mcap`)."
        suffix = Path(filepath).suffix.lower()
        if suffix == ".parquet":
            self.to_parquet(filepath)
            return
        if suffix == ".mcap":
            self.to_mcap(filepath)
            return
        raise ValueError(f"Unsupported file suffix `{suffix}`. Expected one of: `.parquet`, `.mcap`.")

    def apply_projections(self):
        """
        Apply projection transformations to the recording's moving object data based on the provided projection metadata
        and the map's projection. This method updates the `x`, `y`, and `z` columns of the recording's DataFrame
        according to the specified projections and transforms the coordinates to the target CRS if necessary.
        The original coordinates before applying projections are stored in `x_original`, `y_original`, and `z_original`
        columns to preserve the original pose information for export or reference.
        """
        if self._df.height == 0:
            return self

        source_proj_string = self.projections.get("proj_string")
        if source_proj_string is None:
            self.map.parse()
            source_proj_string = getattr(self.map, "proj_string", None)

        if source_proj_string is None:
            raise ValueError("No proj_string information available on the recording or attached map.")

        frame_projections: list[dict[str, typing.Any]] = []
        for ts, offset in self.projections.items():
            if ts in (None, "proj_string"):
                continue
            ox, oy, oz, oyaw = self._offset_components(offset)
            frame_projections.append(
                dict(
                    total_nanos=int(ts),
                    offset_x=ox,
                    offset_y=oy,
                    offset_z=oz,
                    offset_yaw=oyaw,
                )
            )

        df = self._df

        default_offset = self.projections.get(None)
        dox, doy, doz, doyaw = self._offset_components(default_offset)

        if frame_projections:
            proj_df = pl.DataFrame(
                frame_projections,
                schema={
                    "total_nanos": polars_schema["total_nanos"],
                    "offset_x": pl.Float64,
                    "offset_y": pl.Float64,
                    "offset_z": pl.Float64,
                    "offset_yaw": pl.Float64,
                },
            )
            df = df.join(proj_df, on="total_nanos", how="left")
            df = df.with_columns(
                pl.lit(source_proj_string).alias("proj_string"),
                pl.col("offset_x").fill_null(dox).alias("offset_x"),
                pl.col("offset_y").fill_null(doy).alias("offset_y"),
                pl.col("offset_z").fill_null(doz).alias("offset_z"),
                pl.col("offset_yaw").fill_null(doyaw).alias("offset_yaw"),
            )

        else:
            df = df.with_columns(
                pl.lit(source_proj_string).alias("proj_string"),
                pl.lit(dox).alias("offset_x"),
                pl.lit(doy).alias("offset_y"),
                pl.lit(doz).alias("offset_z"),
                pl.lit(doyaw).alias("offset_yaw"),
            )
        source_crs = pyproj.CRS.from_string(source_proj_string)

        if df.select(pl.col("proj_string").is_null().any()).item():
            raise ValueError("Some rows do not have a projection string assigned.")

        # Store original values before applying offsets, when it is the first projection
        if not any(col in df.columns for col in ["x_original", "y_original", "z_original"]):
            df = df.with_columns(
                pl.col("x").alias("x_original"),
                pl.col("y").alias("y_original"),
                pl.col("z").alias("z_original"),
            )

        # Update main columns with offset values
        df = df.with_columns(
            (
                pl.col("x") * pl.col("offset_yaw").cos() - pl.col("y") * pl.col("offset_yaw").sin() + pl.col("offset_x")
            ).alias("x"),
            (
                pl.col("x") * pl.col("offset_yaw").sin() + pl.col("y") * pl.col("offset_yaw").cos() + pl.col("offset_y")
            ).alias("y"),
            (pl.col("z") + pl.col("offset_z")).alias("z"),
        )

        self.map.parse()
        target_crs = self.map.projection
        if not target_crs:
            raise ValueError("Map does not have a valid projection defined.")

        # Apply 2D proj string transformation
        transformer = pyproj.Transformer.from_crs(source_crs, target_crs)
        x_tgt, y_tgt = transformer.transform(df["x"].to_numpy(), df["y"].to_numpy())
        df = df.with_columns(pl.Series(name="x", values=x_tgt), pl.Series(name="y", values=y_tgt))

        # From map world to map local
        if self.map.proj_offset:
            m_ox, m_oy, m_oz, m_oyaw = self._offset_components(self.map.proj_offset)
            df = df.with_columns(
                ((pl.col("x") - m_ox) * np.cos(m_oyaw) + (pl.col("y") - m_oy) * np.sin(m_oyaw)).alias("x"),
                ((pl.col("y") - m_oy) * np.cos(m_oyaw) - (pl.col("x") - m_ox) * np.sin(m_oyaw)).alias("y"),
                (pl.col("z") - m_oz).alias("z"),
            )

        df = bbx_to_polygon(df)

        # Remove temporary projection columns
        df = df.drop("proj_string", "offset_x", "offset_y", "offset_z", "offset_yaw")

        self._df = df
        return self

    def interpolate(self, new_nanos: list[int] | None = None, hz: float | None = None):
        "Interpolate the recording to new timestamps or a given frequency."
        df = self._df
        nanos_min, nanos_max, frame_min, frame_max = df.select(
            nanos_min=pl.col("total_nanos").min(),
            nanos_max=pl.col("total_nanos").max(),
            frame_min=pl.col("frame").min(),
            frame_max=pl.col("frame").max(),
        ).row(0)
        if new_nanos is None:
            if hz is None:
                new_nanos = np.linspace(nanos_min, nanos_max, frame_max - frame_min, dtype=int)
            else:
                step = 1e9 / hz
                new_nanos = np.arange(start=nanos_min, stop=nanos_max + 1, step=step, dtype=int)
        else:
            new_nanos = np.array(new_nanos)
        new_dfs = []
        for [idx], track_df in df.group_by("idx"):
            track_data = {}
            track_new_nanos = new_nanos[
                np.logical_and(
                    track_df["total_nanos"].min() <= new_nanos,
                    track_df["total_nanos"].max() >= new_nanos,
                )
            ]
            for c in [
                "x",
                "y",
                "z",
                "vel_x",
                "vel_y",
                "vel_z",
                "acc_x",
                "acc_y",
                "acc_z",
                "length",
                "width",
                "height",
            ]:
                track_data[c] = np.interp(track_new_nanos, track_df["total_nanos"], track_df[c])
            for c in ["type", "subtype", "role"]:
                track_data[c] = nearest_interp(
                    track_new_nanos,
                    track_df["total_nanos"].to_numpy(),
                    track_df[c].to_numpy(),
                )
            for c in ["roll", "pitch", "yaw"]:
                # Unwrap angles to handle discontinuities, then interpolate, then wrap back to [-π, π]
                unwrapped_angles = np.unwrap(track_df[c])
                interpolated = np.interp(track_new_nanos, track_df["total_nanos"], unwrapped_angles)
                track_data[c] = np.mod(interpolated + np.pi, 2 * np.pi) - np.pi
            new_track_df = pl.DataFrame(track_data)
            new_track_df = new_track_df.with_columns(
                pl.Series(
                    name="idx",
                    values=np.ones_like(track_new_nanos) * idx,
                    dtype=polars_schema["idx"],
                ),
                pl.Series(
                    name="total_nanos",
                    values=track_new_nanos,
                    dtype=polars_schema["total_nanos"],
                ),
            )
            new_dfs.append(new_track_df)
        new_df = pl.concat(new_dfs)
        return self.__init__(df=new_df, map=self.map, host_vehicle_idx=self.host_vehicle_idx)

    def _create_legend(self, ax):
        handles, labels = ax.get_legend_handles_labels()
        host_label = f"{self.host_vehicle_idx} - HV"

        def sort_key(item):
            label = item[1]
            if label == host_label:
                return (-1, -1)
            try:
                return (0, int(label))
            except ValueError:
                return (0, float("inf"))  # non-numeric labels go last

        items = sorted(zip(handles, labels), key=sort_key)
        handles, labels = zip(*items)
        ax.legend(handles, labels, loc="center left", bbox_to_anchor=(1, 0.5))
        return ax

    def plot(self, ax=None, legend=False, mvs_plt_type: str = "scatter") -> plt.Axes:
        "Generate a static plot of the recording using Matplotlib. Plots the map (if available), moving objects, and traffic light states."
        if ax is None:
            fig, ax = plt.subplots(1, 1)
            ax.set_aspect(1)
        if self.map:
            self.map.plot(ax)
        self.plot_mvs(ax=ax, mvs_plt_type=mvs_plt_type)
        self.plot_tl(ax=ax)
        if legend:
            ax = self._create_legend(ax)
        return ax

    def plot_mvs(self, ax=None, legend=False, mvs_plt_type: str = "scatter"):
        "Generate a static plot of the moving objects in the recording using Matplotlib."
        if ax is None:
            fig, ax = plt.subplots(1, 1)
            ax.set_aspect(1)
        plot_fn = {"scatter": ax.scatter, "plot": ax.plot}.get(mvs_plt_type)
        if plot_fn is None:
            raise ValueError("`mvs_plt_type` must be one of: 'scatter', 'plot'.")

        plot_df = self._df["idx", "x", "y"]
        base_kwargs = {"alpha": 0.5}
        for [idx], mv in plot_df.group_by("idx"):
            if idx == self.host_vehicle_idx:
                ax.plot(*mv["x", "y"], c="red", label=f"{idx} - HV")
                continue
            plot_fn(*mv["x", "y"], label=str(idx), **base_kwargs)

        if legend:
            ax = self._create_legend(ax)
        return ax

    def plot_tl(self, ax=None):
        "Generate a static plot of the traffic lights in the recording using Matplotlib."
        if ax is None:
            fig, ax = plt.subplots(1, 1)
            ax.set_aspect(1)
        tl_dict = {}
        for tl_states in self.traffic_light_states:
            for tl in self.traffic_light_states[tl_states]:
                if tl.id.value not in tl_dict.keys():
                    tl_dict[tl.id.value] = tl

        for tl in tl_dict:
            try:
                x = tl_dict[tl].base.position.x
                y = tl_dict[tl].base.position.y
                ax.plot(
                    x,
                    y,
                    marker="s",
                    label=f"Traffic Light {tl_dict[tl].id.value}",
                    c="blue",
                    alpha=0.7,
                    markersize=2,
                )
            except AttributeError as e:
                print(f"Warning: Skipping traffic light {tl.id.value} due to missing position data: {e}")
                continue
        return ax

    def plot_frame(self, frame: int, ax=None):
        "Generate a static plot of a specific frame in the recording using Matplotlib."
        ax = self.plot(ax=ax)
        self.plot_mv_frame(ax, frame=frame)
        return ax

    def plot_mv_frame(self, ax: plt.Axes, frame: int):
        polys = self._df.filter(pl.col("frame") == frame)["polygon"]
        for p in polys:
            ax.add_patch(PltPolygon(p.exterior.coords, fc="red"))

    def plot_altair(
        self,
        start_frame: int = 0,
        end_frame: int = -1,
        plot_map: bool = True,
        plot_map_polys: bool = True,
        metric_column: str | None = None,
        plot_wedges: bool = True,
        idx=None,
        height: float | None = None,
        width: float | None = None,
    ) -> alt.Chart:
        "Generate an interactive plot of the recording using Altair."
        if end_frame != -1:
            df = self._df.filter(pl.col("frame") < end_frame, pl.col("frame") >= start_frame)
        else:
            df = self._df.filter(pl.col("frame") >= start_frame)

        [frame_min], [frame_max] = df.select(
            pl.col("frame").min().alias("min"),
            pl.col("frame").max().alias("max"),
        )[0]
        slider = alt.binding_range(min=frame_min, max=frame_max, step=1, name="frame")
        op_var = alt.param(value=0, bind=slider)

        df = df.with_columns(
            pl.concat_str(
                pl.col("type").map_elements(lambda x: betterosi.MovingObjectType(x).name, return_dtype=pl.String),
                pl.col("subtype").map_elements(
                    lambda x: betterosi.MovingObjectVehicleClassificationType(x).name,
                    return_dtype=pl.String,
                ),
                separator="-",
            ).alias("type")
        )
        buffer = pl.col("length").max()
        xmin, xmax, ymin, ymax = df.select(
            (pl.col("x").min() - buffer).alias("xmin"),
            (pl.col("x").max() + buffer).alias("xmax"),
            (pl.col("y").min() - buffer).alias("ymin"),
            (pl.col("y").max() + buffer).alias("ymax"),
        ).row(0)
        pov_df = pl.DataFrame({"polygon": [shapely.Polygon([[xmax, ymax], [xmax, ymin], [xmin, ymin], [xmin, ymax]])]})
        pov_df = pov_df.select(geometry=st.from_shapely("polygon"))
        pov = alt.Chart({"values": pov_df.st.to_dicts()}).mark_geoshape(fillOpacity=0, filled=False, opacity=0)

        plots = [pov]
        if plot_map and self.map is not None:
            plots.append(self.map.plot_altair(recording=self, plot_polys=plot_map_polys))

        mv_dict = {"values": df["geometry", "idx", "frame", "type"].st.to_dicts()}
        plots.append(
            alt.Chart(mv_dict)
            .mark_geoshape()
            .encode(
                color=(
                    alt.when(alt.FieldEqualPredicate(equal=self.host_vehicle_idx or -1, field="properties.idx"))
                    .then(alt.value("red"))
                    .when(alt.FieldEqualPredicate(equal=-1 if idx is None else idx, field="properties.idx"))
                    .then(alt.value("red"))
                    .otherwise(alt.value("blue"))
                ),
                tooltip=["properties.idx:N", "properties.frame:N", "properties.type:O"],
            )
            .transform_filter(alt.FieldEqualPredicate(field="properties.frame", equal=op_var))
        )
        if plot_wedges:
            wedges_df = df["idx", "frame", "type", "x", "y", "yaw", "length"].with_columns(
                pl.col("yaw").degrees().alias("deg"),
                (pl.col("length") / 4).alias("size"),
            )
            plots.append(
                alt.Chart(wedges_df)
                .mark_point(shape="wedge", color="white", strokeWidth=2)
                .encode(
                    alt.Longitude("x:Q"),
                    alt.Latitude("y:Q"),
                    alt.Angle("deg").scale(domain=[180, -180], range=[-90, 270]),
                    alt.Size("size", legend=None),
                    tooltip=["idx:N", "frame:N", "type:O"],
                )
                .transform_filter(alt.FieldEqualPredicate(field="frame", equal=op_var))
            )

        view = (
            alt.layer(*plots)
            .properties(
                title="Map",
                **({"height": height} if height is not None else {}),
                **({"width": width} if width is not None else {}),
            )
            .project("identity", reflectY=True)
        )

        if metric_column is not None and idx is not None:
            metric = (
                df["idx", metric_column, "frame"]
                .filter(idx=idx)
                .plot.line(x="frame", y=metric_column, color=alt.value("red"))
                .properties(title=f"{metric_column} of object {idx}")
            )
            vertline = (
                alt.Chart()
                .mark_rule()
                .encode(
                    x=alt.datum(
                        op_var,
                        type="quantitative",
                        scale=alt.Scale(domain=[frame_min, frame_max]),
                    )
                )
            )
            view = view | (metric + vertline)
        return view.add_params(op_var)

    def create_mapsegments(self):
        if isinstance(self.map, MapOsiCenterline):
            self.mapsegment = MapOsiCenterlineSegmentation(self)
            self.mapsegment.init_intersections()

__init__(df, map=None, projections=None, host_vehicle_idx=None, validate=False, traffic_light_states=None)

Initialize a Recording instance.

Source code in omega_prime/recording.py

def __init__(
    self,
    df,
    map=None,
    projections=None,
    host_vehicle_idx: int | None = None,
    validate=False,
    traffic_light_states: dict | None = None,
):
    "Initialize a Recording instance."
    df = self._ensure_polars_dataframe(df)
    if "total_nanos" not in df.columns:
        raise ValueError("df must contain column `total_nanos`.")
    nanos2frame, mapping = self._build_frame_mapping(df)
    df = self._attach_frame_column(df, mapping)
    df = self._ensure_polars_dataframe(df)
    if validate:
        recording_moving_object_schema.validate(df, lazy=True)

    super().__init__()
    self.nanos2frame = nanos2frame

    df = self._ensure_motion_norm_columns(df)
    self.projections = self._validate_projections_schema(projections)
    self.traffic_light_states = traffic_light_states if traffic_light_states is not None else {}

    df = bbx_to_polygon(df)

    self._df = df
    self.map = map
    self._moving_objects = None
    self.host_vehicle_idx = host_vehicle_idx
    self.mapsegment = None

apply_projections()

Apply projection transformations to the recording's moving object data based on the provided projection metadata and the map's projection. This method updates the x, y, and z columns of the recording's DataFrame according to the specified projections and transforms the coordinates to the target CRS if necessary. The original coordinates before applying projections are stored in x_original, y_original, and z_original columns to preserve the original pose information for export or reference.

Source code in omega_prime/recording.py

def apply_projections(self):
    """
    Apply projection transformations to the recording's moving object data based on the provided projection metadata
    and the map's projection. This method updates the `x`, `y`, and `z` columns of the recording's DataFrame
    according to the specified projections and transforms the coordinates to the target CRS if necessary.
    The original coordinates before applying projections are stored in `x_original`, `y_original`, and `z_original`
    columns to preserve the original pose information for export or reference.
    """
    if self._df.height == 0:
        return self

    source_proj_string = self.projections.get("proj_string")
    if source_proj_string is None:
        self.map.parse()
        source_proj_string = getattr(self.map, "proj_string", None)

    if source_proj_string is None:
        raise ValueError("No proj_string information available on the recording or attached map.")

    frame_projections: list[dict[str, typing.Any]] = []
    for ts, offset in self.projections.items():
        if ts in (None, "proj_string"):
            continue
        ox, oy, oz, oyaw = self._offset_components(offset)
        frame_projections.append(
            dict(
                total_nanos=int(ts),
                offset_x=ox,
                offset_y=oy,
                offset_z=oz,
                offset_yaw=oyaw,
            )
        )

    df = self._df

    default_offset = self.projections.get(None)
    dox, doy, doz, doyaw = self._offset_components(default_offset)

    if frame_projections:
        proj_df = pl.DataFrame(
            frame_projections,
            schema={
                "total_nanos": polars_schema["total_nanos"],
                "offset_x": pl.Float64,
                "offset_y": pl.Float64,
                "offset_z": pl.Float64,
                "offset_yaw": pl.Float64,
            },
        )
        df = df.join(proj_df, on="total_nanos", how="left")
        df = df.with_columns(
            pl.lit(source_proj_string).alias("proj_string"),
            pl.col("offset_x").fill_null(dox).alias("offset_x"),
            pl.col("offset_y").fill_null(doy).alias("offset_y"),
            pl.col("offset_z").fill_null(doz).alias("offset_z"),
            pl.col("offset_yaw").fill_null(doyaw).alias("offset_yaw"),
        )

    else:
        df = df.with_columns(
            pl.lit(source_proj_string).alias("proj_string"),
            pl.lit(dox).alias("offset_x"),
            pl.lit(doy).alias("offset_y"),
            pl.lit(doz).alias("offset_z"),
            pl.lit(doyaw).alias("offset_yaw"),
        )
    source_crs = pyproj.CRS.from_string(source_proj_string)

    if df.select(pl.col("proj_string").is_null().any()).item():
        raise ValueError("Some rows do not have a projection string assigned.")

    # Store original values before applying offsets, when it is the first projection
    if not any(col in df.columns for col in ["x_original", "y_original", "z_original"]):
        df = df.with_columns(
            pl.col("x").alias("x_original"),
            pl.col("y").alias("y_original"),
            pl.col("z").alias("z_original"),
        )

    # Update main columns with offset values
    df = df.with_columns(
        (
            pl.col("x") * pl.col("offset_yaw").cos() - pl.col("y") * pl.col("offset_yaw").sin() + pl.col("offset_x")
        ).alias("x"),
        (
            pl.col("x") * pl.col("offset_yaw").sin() + pl.col("y") * pl.col("offset_yaw").cos() + pl.col("offset_y")
        ).alias("y"),
        (pl.col("z") + pl.col("offset_z")).alias("z"),
    )

    self.map.parse()
    target_crs = self.map.projection
    if not target_crs:
        raise ValueError("Map does not have a valid projection defined.")

    # Apply 2D proj string transformation
    transformer = pyproj.Transformer.from_crs(source_crs, target_crs)
    x_tgt, y_tgt = transformer.transform(df["x"].to_numpy(), df["y"].to_numpy())
    df = df.with_columns(pl.Series(name="x", values=x_tgt), pl.Series(name="y", values=y_tgt))

    # From map world to map local
    if self.map.proj_offset:
        m_ox, m_oy, m_oz, m_oyaw = self._offset_components(self.map.proj_offset)
        df = df.with_columns(
            ((pl.col("x") - m_ox) * np.cos(m_oyaw) + (pl.col("y") - m_oy) * np.sin(m_oyaw)).alias("x"),
            ((pl.col("y") - m_oy) * np.cos(m_oyaw) - (pl.col("x") - m_ox) * np.sin(m_oyaw)).alias("y"),
            (pl.col("z") - m_oz).alias("z"),
        )

    df = bbx_to_polygon(df)

    # Remove temporary projection columns
    df = df.drop("proj_string", "offset_x", "offset_y", "offset_z", "offset_yaw")

    self._df = df
    return self

from_file(filepath, map_path=None, validate=False, parse_map=False, step_size=0.01, apply_proj=True, **kwargs) classmethod

Load a Recording from a file. Supports .parquet, .osi and .mcap files.

Parameters:

Name	Type	Description	Default
filepath	str	Path to the input file.	required
map_path	str | None	Optional path to a map file. If None, the map will be loaded from the recording if available.	None
validate	bool	Whether to validate the data against the schema.	False
parse_map	bool	Whether to create python objects from the map data or just load it.	False
step_size	float	Step size for map parsing, if applicable (Used for ASAM OpenDRIVE).	0.01
apply_proj	bool	Whether to apply projection transformations to the recording's moving object data.	True

Returns:

Name	Type	Description
Recording	Recording	The loaded Recording object.

Source code in omega_prime/recording.py

@classmethod
def from_file(
    cls,
    filepath,
    map_path: str | None = None,
    validate: bool = False,
    parse_map: bool = False,
    step_size: float = 0.01,
    apply_proj: bool = True,
    **kwargs,
) -> "Recording":
    """Load a Recording from a file. Supports `.parquet`, `.osi` and `.mcap` files.

    Parameters:
        filepath (str): Path to the input file.
        map_path (str | None): Optional path to a map file. If None, the map will be loaded from the recording if available.
        validate (bool): Whether to validate the data against the schema.
        parse_map (bool): Whether to create python objects from the map data or just load it.
        step_size (float): Step size for map parsing, if applicable (Used for ASAM OpenDRIVE).
        apply_proj (bool): Whether to apply projection transformations to the recording's moving object data.

    Returns:
        Recording (Recording): The loaded Recording object.
    """
    if filepath is None and map_path is None:
        raise ValueError("Either `filepath` or `map_path` must be provided.")

    if filepath is not None and Path(filepath).suffix == ".parquet":
        r = cls.from_parquet(filepath, parse_map=parse_map, validate=validate, step_size=step_size)
    elif filepath is not None:
        gts = betterosi.read(filepath, return_ground_truth=True, mcap_return_betterosi=True)
        r = cls.from_osi_gts(gts, validate=validate)
    if map_path is None and r.map is not None:
        return r

    map_path = Path(map_path if map_path is not None else filepath)
    map_parsing = {}
    map = None
    for MC in MAP_CLASSES:
        if map_path.suffix in MC._supported_file_suffixes:
            try:
                map = MC.from_file(map_path, parse_map=parse_map, **kwargs)
            except Exception as e:
                map_parsing[MC.__name__] = str(e)
            else:
                break
    if map is not None:
        r.map = map
    elif r.map is None:
        warn(f"No map could be found: {map_parsing}")

    if r.projections and apply_proj:
        try:
            r.apply_projections()
        except Exception:
            warn("Failed to apply projections.")
    return r

interpolate(new_nanos=None, hz=None)

Interpolate the recording to new timestamps or a given frequency.

Source code in omega_prime/recording.py

def interpolate(self, new_nanos: list[int] | None = None, hz: float | None = None):
    "Interpolate the recording to new timestamps or a given frequency."
    df = self._df
    nanos_min, nanos_max, frame_min, frame_max = df.select(
        nanos_min=pl.col("total_nanos").min(),
        nanos_max=pl.col("total_nanos").max(),
        frame_min=pl.col("frame").min(),
        frame_max=pl.col("frame").max(),
    ).row(0)
    if new_nanos is None:
        if hz is None:
            new_nanos = np.linspace(nanos_min, nanos_max, frame_max - frame_min, dtype=int)
        else:
            step = 1e9 / hz
            new_nanos = np.arange(start=nanos_min, stop=nanos_max + 1, step=step, dtype=int)
    else:
        new_nanos = np.array(new_nanos)
    new_dfs = []
    for [idx], track_df in df.group_by("idx"):
        track_data = {}
        track_new_nanos = new_nanos[
            np.logical_and(
                track_df["total_nanos"].min() <= new_nanos,
                track_df["total_nanos"].max() >= new_nanos,
            )
        ]
        for c in [
            "x",
            "y",
            "z",
            "vel_x",
            "vel_y",
            "vel_z",
            "acc_x",
            "acc_y",
            "acc_z",
            "length",
            "width",
            "height",
        ]:
            track_data[c] = np.interp(track_new_nanos, track_df["total_nanos"], track_df[c])
        for c in ["type", "subtype", "role"]:
            track_data[c] = nearest_interp(
                track_new_nanos,
                track_df["total_nanos"].to_numpy(),
                track_df[c].to_numpy(),
            )
        for c in ["roll", "pitch", "yaw"]:
            # Unwrap angles to handle discontinuities, then interpolate, then wrap back to [-π, π]
            unwrapped_angles = np.unwrap(track_df[c])
            interpolated = np.interp(track_new_nanos, track_df["total_nanos"], unwrapped_angles)
            track_data[c] = np.mod(interpolated + np.pi, 2 * np.pi) - np.pi
        new_track_df = pl.DataFrame(track_data)
        new_track_df = new_track_df.with_columns(
            pl.Series(
                name="idx",
                values=np.ones_like(track_new_nanos) * idx,
                dtype=polars_schema["idx"],
            ),
            pl.Series(
                name="total_nanos",
                values=track_new_nanos,
                dtype=polars_schema["total_nanos"],
            ),
        )
        new_dfs.append(new_track_df)
    new_df = pl.concat(new_dfs)
    return self.__init__(df=new_df, map=self.map, host_vehicle_idx=self.host_vehicle_idx)

plot(ax=None, legend=False, mvs_plt_type='scatter')

Generate a static plot of the recording using Matplotlib. Plots the map (if available), moving objects, and traffic light states.

Source code in omega_prime/recording.py

def plot(self, ax=None, legend=False, mvs_plt_type: str = "scatter") -> plt.Axes:
    "Generate a static plot of the recording using Matplotlib. Plots the map (if available), moving objects, and traffic light states."
    if ax is None:
        fig, ax = plt.subplots(1, 1)
        ax.set_aspect(1)
    if self.map:
        self.map.plot(ax)
    self.plot_mvs(ax=ax, mvs_plt_type=mvs_plt_type)
    self.plot_tl(ax=ax)
    if legend:
        ax = self._create_legend(ax)
    return ax

plot_altair(start_frame=0, end_frame=-1, plot_map=True, plot_map_polys=True, metric_column=None, plot_wedges=True, idx=None, height=None, width=None)

Generate an interactive plot of the recording using Altair.

Source code in omega_prime/recording.py

def plot_altair(
    self,
    start_frame: int = 0,
    end_frame: int = -1,
    plot_map: bool = True,
    plot_map_polys: bool = True,
    metric_column: str | None = None,
    plot_wedges: bool = True,
    idx=None,
    height: float | None = None,
    width: float | None = None,
) -> alt.Chart:
    "Generate an interactive plot of the recording using Altair."
    if end_frame != -1:
        df = self._df.filter(pl.col("frame") < end_frame, pl.col("frame") >= start_frame)
    else:
        df = self._df.filter(pl.col("frame") >= start_frame)

    [frame_min], [frame_max] = df.select(
        pl.col("frame").min().alias("min"),
        pl.col("frame").max().alias("max"),
    )[0]
    slider = alt.binding_range(min=frame_min, max=frame_max, step=1, name="frame")
    op_var = alt.param(value=0, bind=slider)

    df = df.with_columns(
        pl.concat_str(
            pl.col("type").map_elements(lambda x: betterosi.MovingObjectType(x).name, return_dtype=pl.String),
            pl.col("subtype").map_elements(
                lambda x: betterosi.MovingObjectVehicleClassificationType(x).name,
                return_dtype=pl.String,
            ),
            separator="-",
        ).alias("type")
    )
    buffer = pl.col("length").max()
    xmin, xmax, ymin, ymax = df.select(
        (pl.col("x").min() - buffer).alias("xmin"),
        (pl.col("x").max() + buffer).alias("xmax"),
        (pl.col("y").min() - buffer).alias("ymin"),
        (pl.col("y").max() + buffer).alias("ymax"),
    ).row(0)
    pov_df = pl.DataFrame({"polygon": [shapely.Polygon([[xmax, ymax], [xmax, ymin], [xmin, ymin], [xmin, ymax]])]})
    pov_df = pov_df.select(geometry=st.from_shapely("polygon"))
    pov = alt.Chart({"values": pov_df.st.to_dicts()}).mark_geoshape(fillOpacity=0, filled=False, opacity=0)

    plots = [pov]
    if plot_map and self.map is not None:
        plots.append(self.map.plot_altair(recording=self, plot_polys=plot_map_polys))

    mv_dict = {"values": df["geometry", "idx", "frame", "type"].st.to_dicts()}
    plots.append(
        alt.Chart(mv_dict)
        .mark_geoshape()
        .encode(
            color=(
                alt.when(alt.FieldEqualPredicate(equal=self.host_vehicle_idx or -1, field="properties.idx"))
                .then(alt.value("red"))
                .when(alt.FieldEqualPredicate(equal=-1 if idx is None else idx, field="properties.idx"))
                .then(alt.value("red"))
                .otherwise(alt.value("blue"))
            ),
            tooltip=["properties.idx:N", "properties.frame:N", "properties.type:O"],
        )
        .transform_filter(alt.FieldEqualPredicate(field="properties.frame", equal=op_var))
    )
    if plot_wedges:
        wedges_df = df["idx", "frame", "type", "x", "y", "yaw", "length"].with_columns(
            pl.col("yaw").degrees().alias("deg"),
            (pl.col("length") / 4).alias("size"),
        )
        plots.append(
            alt.Chart(wedges_df)
            .mark_point(shape="wedge", color="white", strokeWidth=2)
            .encode(
                alt.Longitude("x:Q"),
                alt.Latitude("y:Q"),
                alt.Angle("deg").scale(domain=[180, -180], range=[-90, 270]),
                alt.Size("size", legend=None),
                tooltip=["idx:N", "frame:N", "type:O"],
            )
            .transform_filter(alt.FieldEqualPredicate(field="frame", equal=op_var))
        )

    view = (
        alt.layer(*plots)
        .properties(
            title="Map",
            **({"height": height} if height is not None else {}),
            **({"width": width} if width is not None else {}),
        )
        .project("identity", reflectY=True)
    )

    if metric_column is not None and idx is not None:
        metric = (
            df["idx", metric_column, "frame"]
            .filter(idx=idx)
            .plot.line(x="frame", y=metric_column, color=alt.value("red"))
            .properties(title=f"{metric_column} of object {idx}")
        )
        vertline = (
            alt.Chart()
            .mark_rule()
            .encode(
                x=alt.datum(
                    op_var,
                    type="quantitative",
                    scale=alt.Scale(domain=[frame_min, frame_max]),
                )
            )
        )
        view = view | (metric + vertline)
    return view.add_params(op_var)

plot_frame(frame, ax=None)

Generate a static plot of a specific frame in the recording using Matplotlib.

Source code in omega_prime/recording.py

def plot_frame(self, frame: int, ax=None):
    "Generate a static plot of a specific frame in the recording using Matplotlib."
    ax = self.plot(ax=ax)
    self.plot_mv_frame(ax, frame=frame)
    return ax

plot_mvs(ax=None, legend=False, mvs_plt_type='scatter')

Generate a static plot of the moving objects in the recording using Matplotlib.

Source code in omega_prime/recording.py

def plot_mvs(self, ax=None, legend=False, mvs_plt_type: str = "scatter"):
    "Generate a static plot of the moving objects in the recording using Matplotlib."
    if ax is None:
        fig, ax = plt.subplots(1, 1)
        ax.set_aspect(1)
    plot_fn = {"scatter": ax.scatter, "plot": ax.plot}.get(mvs_plt_type)
    if plot_fn is None:
        raise ValueError("`mvs_plt_type` must be one of: 'scatter', 'plot'.")

    plot_df = self._df["idx", "x", "y"]
    base_kwargs = {"alpha": 0.5}
    for [idx], mv in plot_df.group_by("idx"):
        if idx == self.host_vehicle_idx:
            ax.plot(*mv["x", "y"], c="red", label=f"{idx} - HV")
            continue
        plot_fn(*mv["x", "y"], label=str(idx), **base_kwargs)

    if legend:
        ax = self._create_legend(ax)
    return ax

plot_tl(ax=None)

Generate a static plot of the traffic lights in the recording using Matplotlib.

Source code in omega_prime/recording.py

def plot_tl(self, ax=None):
    "Generate a static plot of the traffic lights in the recording using Matplotlib."
    if ax is None:
        fig, ax = plt.subplots(1, 1)
        ax.set_aspect(1)
    tl_dict = {}
    for tl_states in self.traffic_light_states:
        for tl in self.traffic_light_states[tl_states]:
            if tl.id.value not in tl_dict.keys():
                tl_dict[tl.id.value] = tl

    for tl in tl_dict:
        try:
            x = tl_dict[tl].base.position.x
            y = tl_dict[tl].base.position.y
            ax.plot(
                x,
                y,
                marker="s",
                label=f"Traffic Light {tl_dict[tl].id.value}",
                c="blue",
                alpha=0.7,
                markersize=2,
            )
        except AttributeError as e:
            print(f"Warning: Skipping traffic light {tl.id.value} due to missing position data: {e}")
            continue
    return ax

to_file(filepath)

Store Recording to a file based on its suffix (.parquet, .mcap).

Source code in omega_prime/recording.py

def to_file(self, filepath):
    "Store Recording to a file based on its suffix (`.parquet`, `.mcap`)."
    suffix = Path(filepath).suffix.lower()
    if suffix == ".parquet":
        self.to_parquet(filepath)
        return
    if suffix == ".mcap":
        self.to_mcap(filepath)
        return
    raise ValueError(f"Unsupported file suffix `{suffix}`. Expected one of: `.parquet`, `.mcap`.")

to_mcap(filepath)

Store Recording as an MCAP file.

Source code in omega_prime/recording.py

def to_mcap(self, filepath):
    "Store Recording as an MCAP file."
    if Path(filepath).suffix != ".mcap":
        raise ValueError()
    with betterosi.Writer(filepath) as w:
        for gt in self.to_osi_gts():
            w.add(gt)
        if isinstance(self.map, MapOdr):
            w.add(self.map.to_osi(), topic="ground_truth_map", log_time=0)
        elif (
            self.map is not None and not isinstance(self.map, MapOsi) and not isinstance(self.map, MapOsiCenterline)
        ):
            warn(f"The map {self.map} could not be saved to `mcap`")

to_parquet(filename)

Store Recording as a Parquet file.

Source code in omega_prime/recording.py

def to_parquet(self, filename):
    "Store Recording as a Parquet file."
    metadata = {}
    if self.host_vehicle_idx is not None:
        metadata[b"host_vehicle_idx"] = str(self.host_vehicle_idx).encode()
    proj_meta = {}
    encoded_projections = self._encode_projections(self.projections)
    if encoded_projections:
        proj_meta[b"projections_json"] = encoded_projections
    df_export = self._df_with_original_pose_for_export()
    to_drop = ["frame"]
    optional_cols = [
        "polygon",
        "global_lat",
        "global_lon",
        "global_alt",
        "global_yaw",
        "proj_string",
        "x_original",
        "y_original",
        "z_original",
        "yaw_original",
    ]
    to_drop.extend([c for c in optional_cols if c in df_export.columns])
    t = pyarrow.table(df_export.drop(*to_drop))
    map_meta = self.map._to_binary_json() if self.map is not None else {}

    t = t.cast(t.schema.with_metadata(metadata | proj_meta | map_meta))
    pq.write_table(t, filename)

Map dataclass

Base class for Map representations

Source code in omega_prime/map.py

@dataclass(repr=False)
class Map:
    """Base class for Map representations"""

    lane_boundaries: dict[Any, LaneBoundary]
    lanes: dict[Any:Lane]

    _supported_file_suffixes = [".osi", ".mcap"]
    _binary_json_identifier = b"osi"

    def plot(self, ax: plt.Axes | None = None):
        if ax is None:
            fig, ax = plt.subplots(1, 1)
            ax.set_aspect(1)
        for l in self.lanes.values():
            l.plot(ax)
        for b in self.lane_boundaries.values():
            b.plot(ax)

    @classmethod
    def from_file(cls, filepath, parse_map=True, **kwargs):
        "Create a Map instance from a file."
        first_gt = next(betterosi.read(filepath, return_ground_truth=True, mcap_return_betterosi=True))
        return cls.create(first_gt, **kwargs)

    def plot_altair(self, recording=None, plot_polys=True):
        arbitrary_lane = next(iter(self.lanes.values()))
        plot_polys = hasattr(arbitrary_lane, "polygon") and arbitrary_lane.polygon is not None and plot_polys

        if not hasattr(self, "_plot_dict"):
            if plot_polys:
                shapely_series = pl.Series(
                    name="shapely", values=[l.polygon.simplify(0.1) for l in self.lanes.values()]
                )
            else:
                shapely_series = pl.Series(
                    name="shapely", values=[l.centerline.simplify(0.1) for l in self.lanes.values()]
                )

            map_df = pl.DataFrame(
                [
                    shapely_series,
                    pl.Series(name="idx", values=[i for i, _ in enumerate(self.lanes.keys())]),
                    pl.Series(name="type", values=[o.type.name for o in self.lanes.values()]),
                    pl.Series(name="subtype", values=[o.subtype.name for o in self.lanes.values()]),
                    pl.Series(name="on_intersection", values=[o.on_intersection for o in self.lanes.values()]),
                ]
            )
            map_df = map_df.with_columns(geometry=st.from_shapely("shapely")).drop("shapely")

            if recording is not None:
                buffer = 5
                [xmin], [xmax], [ymin], [ymax] = recording._df.select(
                    (pl.col("x").min() - buffer).alias("xmin"),
                    (pl.col("x").max() + buffer).alias("xmax"),
                    (pl.col("y").min() - buffer).alias("ymin"),
                    (pl.col("y").max() + buffer).alias("ymax"),
                )[0]

                pov_df = pl.DataFrame(
                    {"polygon": [shapely.Polygon([[xmax, ymax], [xmax, ymin], [xmin, ymin], [xmin, ymax]])]}
                )
                pov_df = pov_df.select(geometry=st.from_shapely("polygon"))
                map_df = map_df.with_columns(
                    pl.col("geometry").st.intersection(pl.lit(pov_df["geometry"])),
                )
            self._plot_dict = {"values": map_df.st.to_dicts()}

        c = (
            alt.Chart(self._plot_dict)
            .mark_geoshape(fillOpacity=0.4, filled=True if plot_polys else False)
            .encode(
                tooltip=[
                    "properties.idx:N",
                    "properties.type:O",
                    "properties.subtype:O",
                    "properties.on_intersection:O",
                ],
                color=(
                    alt.when(alt.FieldEqualPredicate(equal=True, field="properties.on_intersection"))
                    .then(alt.value("black"))
                    .otherwise(alt.value("green"))
                ),
            )
        )
        if recording is None:
            return c.properties(title="Map").project("identity", reflectY=True)
        else:
            return c

    def map_to_centerline_mcap(self, output_mcap_path: Path = None) -> betterosi.GroundTruth:
        """
        Convert an Map to a MapOsiCenterline and save it as an MCAP file if the output path is provided.
        It returns the generated GroundTruth object from the generated MapOsiCenterline.

        Args:
            output_mcap_path: Path where the MCAP file will be saved
        Returns:
            betterosi.GroundTruth: The generated GroundTruth object
        """

        # Create a mapping from XodrLaneId to a simple integer ID
        lane_id_mapping = {}
        for idx, lane_idx in enumerate(self.lanes.keys()):
            lane_id_mapping[lane_idx] = idx

        # Create betterosi.Lane objects for each lane
        osi_lanes = []
        for lane in self.lanes.values():
            if not lane.centerline.is_valid or lane.centerline.is_empty:
                logging.warning(f"Warning: Skipping invalid lane {lane.idx}")
                continue

            # Check for NaN/inf coordinates
            coords = np.array(lane.centerline.coords)
            if not np.isfinite(coords).all():
                logging.warning(f"Warning: Lane {lane.idx} has non-finite coordinates, skipping")
                continue

            if len(coords) < 2:
                logging.warning(f"Warning: Lane {lane.idx} has insufficient points, skipping")
                continue
            # Get centerline coordinates
            centerline_coords = list(shapely.simplify(lane.centerline, 0.1).coords)
            if not len(centerline_coords) > 1:
                centerline_coords = list(lane.centerline.coords)
                if not len(centerline_coords) > 1:
                    # skip lanes with insufficient centerline points
                    logging.warning(f"Warning: Skipping lane {lane.idx} due to insufficient centerline points")
                    continue

            centerline = [betterosi.Vector3D(x=float(x), y=float(y), z=0.0) for x, y in centerline_coords]

            assert len(centerline_coords) > 1
            # Create lane pairing for successor/predecessor relationships
            lane_pairings = []

            # Get all unique combinations of predecessors and successors
            predecessors = [pred_id for pred_id in lane.predecessor_ids if pred_id in lane_id_mapping]
            successors = [succ_id for succ_id in lane.successor_ids if succ_id in lane_id_mapping]

            # If there are no predecessors or successors, create a single pairing with None values
            if predecessors or successors:
                # Create pairings for all combinations
                if not predecessors:
                    predecessors = [None]
                if not successors:
                    successors = [None]

                for pred_id in predecessors:
                    for succ_id in successors:
                        lane_pairings.append(
                            betterosi.LaneClassificationLanePairing(
                                antecessor_lane_id=betterosi.Identifier(value=lane_id_mapping[pred_id])
                                if pred_id is not None
                                else None,
                                successor_lane_id=betterosi.Identifier(value=lane_id_mapping[succ_id])
                                if succ_id is not None
                                else None,
                            )
                        )

            # Create the OSI lane
            osi_lane = betterosi.Lane(
                id=betterosi.Identifier(value=lane_id_mapping[lane.idx]),
                classification=betterosi.LaneClassification(
                    centerline=centerline,
                    centerline_is_driving_direction=True,
                    type=lane.type,
                    subtype=lane.subtype,
                    lane_pairing=lane_pairings,
                ),
            )
            osi_lanes.append(osi_lane)

        # Create a GroundTruth with only the lanes (no moving objects, no lane boundaries)
        ground_truth = betterosi.GroundTruth(
            version=betterosi.InterfaceVersion(
                version_major=3,
                version_minor=7,
                version_patch=0,
            ),
            timestamp=betterosi.Timestamp(
                seconds=0,
                nanos=0,
            ),
            lane=osi_lanes,
        )

        # Save to MCAP file if output path is provided
        if output_mcap_path is None:
            logging.warning("No output path provided for MCAP file")
        else:
            # Convert string to Path if needed
            output_mcap_path = Path(output_mcap_path)

            if output_mcap_path.is_dir():
                output_mcap_path = output_mcap_path / "map_to_centerline.mcap"
            elif not output_mcap_path.suffix == ".mcap":
                logging.warning(f"Output path must be a directory or .mcap file: {output_mcap_path}")
                return ground_truth

            with betterosi.Writer(output_mcap_path) as writer:
                writer.add(ground_truth, topic="ground_truth_map", log_time=0)
            logging.info(f"Successfully saved map with {len(osi_lanes)} lanes to {output_mcap_path}")

        return ground_truth

    def align_predecessor_and_successor_relations(self):
        """
        Ensure that predecessor and successor relationships between lanes are consistent.
        If lane A lists lane B as a successor, then lane B should list lane A as a predecessor, and vice versa.
        """
        for lane in self.lanes.values():
            for succ_id in lane.successor_ids:
                if succ_id in self.lanes:
                    succ_lane = self.lanes[succ_id]
                    if lane.idx not in succ_lane.predecessor_ids:
                        succ_lane.predecessor_ids.append(lane.idx)
            for pred_id in lane.predecessor_ids:
                if pred_id in self.lanes:
                    pred_lane = self.lanes[pred_id]
                    if lane.idx not in pred_lane.successor_ids:
                        pred_lane.successor_ids.append(lane.idx)

    @classmethod
    def create(cls, *args, **kwargs):
        raise NotImplementedError()

    def __post_init__(self):
        self.setup_lanes_and_boundaries()

    def setup_lanes_and_boundaries(self):
        raise NotImplementedError()

    def _to_binary_json(self):
        raise NotImplementedError()

    @classmethod
    def _from_binary_json(cls, d, **kwargs):
        raise NotImplementedError()

align_predecessor_and_successor_relations()

Ensure that predecessor and successor relationships between lanes are consistent. If lane A lists lane B as a successor, then lane B should list lane A as a predecessor, and vice versa.

Source code in omega_prime/map.py

def align_predecessor_and_successor_relations(self):
    """
    Ensure that predecessor and successor relationships between lanes are consistent.
    If lane A lists lane B as a successor, then lane B should list lane A as a predecessor, and vice versa.
    """
    for lane in self.lanes.values():
        for succ_id in lane.successor_ids:
            if succ_id in self.lanes:
                succ_lane = self.lanes[succ_id]
                if lane.idx not in succ_lane.predecessor_ids:
                    succ_lane.predecessor_ids.append(lane.idx)
        for pred_id in lane.predecessor_ids:
            if pred_id in self.lanes:
                pred_lane = self.lanes[pred_id]
                if lane.idx not in pred_lane.successor_ids:
                    pred_lane.successor_ids.append(lane.idx)

from_file(filepath, parse_map=True, **kwargs) classmethod

Create a Map instance from a file.

Source code in omega_prime/map.py

@classmethod
def from_file(cls, filepath, parse_map=True, **kwargs):
    "Create a Map instance from a file."
    first_gt = next(betterosi.read(filepath, return_ground_truth=True, mcap_return_betterosi=True))
    return cls.create(first_gt, **kwargs)

map_to_centerline_mcap(output_mcap_path=None)

Convert an Map to a MapOsiCenterline and save it as an MCAP file if the output path is provided. It returns the generated GroundTruth object from the generated MapOsiCenterline.

Parameters:

Name	Type	Description	Default
output_mcap_path	Path	Path where the MCAP file will be saved	None

Returns: betterosi.GroundTruth: The generated GroundTruth object

Source code in omega_prime/map.py

def map_to_centerline_mcap(self, output_mcap_path: Path = None) -> betterosi.GroundTruth:
    """
    Convert an Map to a MapOsiCenterline and save it as an MCAP file if the output path is provided.
    It returns the generated GroundTruth object from the generated MapOsiCenterline.

    Args:
        output_mcap_path: Path where the MCAP file will be saved
    Returns:
        betterosi.GroundTruth: The generated GroundTruth object
    """

    # Create a mapping from XodrLaneId to a simple integer ID
    lane_id_mapping = {}
    for idx, lane_idx in enumerate(self.lanes.keys()):
        lane_id_mapping[lane_idx] = idx

    # Create betterosi.Lane objects for each lane
    osi_lanes = []
    for lane in self.lanes.values():
        if not lane.centerline.is_valid or lane.centerline.is_empty:
            logging.warning(f"Warning: Skipping invalid lane {lane.idx}")
            continue

        # Check for NaN/inf coordinates
        coords = np.array(lane.centerline.coords)
        if not np.isfinite(coords).all():
            logging.warning(f"Warning: Lane {lane.idx} has non-finite coordinates, skipping")
            continue

        if len(coords) < 2:
            logging.warning(f"Warning: Lane {lane.idx} has insufficient points, skipping")
            continue
        # Get centerline coordinates
        centerline_coords = list(shapely.simplify(lane.centerline, 0.1).coords)
        if not len(centerline_coords) > 1:
            centerline_coords = list(lane.centerline.coords)
            if not len(centerline_coords) > 1:
                # skip lanes with insufficient centerline points
                logging.warning(f"Warning: Skipping lane {lane.idx} due to insufficient centerline points")
                continue

        centerline = [betterosi.Vector3D(x=float(x), y=float(y), z=0.0) for x, y in centerline_coords]

        assert len(centerline_coords) > 1
        # Create lane pairing for successor/predecessor relationships
        lane_pairings = []

        # Get all unique combinations of predecessors and successors
        predecessors = [pred_id for pred_id in lane.predecessor_ids if pred_id in lane_id_mapping]
        successors = [succ_id for succ_id in lane.successor_ids if succ_id in lane_id_mapping]

        # If there are no predecessors or successors, create a single pairing with None values
        if predecessors or successors:
            # Create pairings for all combinations
            if not predecessors:
                predecessors = [None]
            if not successors:
                successors = [None]

            for pred_id in predecessors:
                for succ_id in successors:
                    lane_pairings.append(
                        betterosi.LaneClassificationLanePairing(
                            antecessor_lane_id=betterosi.Identifier(value=lane_id_mapping[pred_id])
                            if pred_id is not None
                            else None,
                            successor_lane_id=betterosi.Identifier(value=lane_id_mapping[succ_id])
                            if succ_id is not None
                            else None,
                        )
                    )

        # Create the OSI lane
        osi_lane = betterosi.Lane(
            id=betterosi.Identifier(value=lane_id_mapping[lane.idx]),
            classification=betterosi.LaneClassification(
                centerline=centerline,
                centerline_is_driving_direction=True,
                type=lane.type,
                subtype=lane.subtype,
                lane_pairing=lane_pairings,
            ),
        )
        osi_lanes.append(osi_lane)

    # Create a GroundTruth with only the lanes (no moving objects, no lane boundaries)
    ground_truth = betterosi.GroundTruth(
        version=betterosi.InterfaceVersion(
            version_major=3,
            version_minor=7,
            version_patch=0,
        ),
        timestamp=betterosi.Timestamp(
            seconds=0,
            nanos=0,
        ),
        lane=osi_lanes,
    )

    # Save to MCAP file if output path is provided
    if output_mcap_path is None:
        logging.warning("No output path provided for MCAP file")
    else:
        # Convert string to Path if needed
        output_mcap_path = Path(output_mcap_path)

        if output_mcap_path.is_dir():
            output_mcap_path = output_mcap_path / "map_to_centerline.mcap"
        elif not output_mcap_path.suffix == ".mcap":
            logging.warning(f"Output path must be a directory or .mcap file: {output_mcap_path}")
            return ground_truth

        with betterosi.Writer(output_mcap_path) as writer:
            writer.add(ground_truth, topic="ground_truth_map", log_time=0)
        logging.info(f"Successfully saved map with {len(osi_lanes)} lanes to {output_mcap_path}")

    return ground_truth

MapOsi dataclass

Bases: Map

Map representation based on ASAM OSI GroundTruth

Source code in omega_prime/map.py

@dataclass(repr=False)
class MapOsi(Map):
    "Map representation based on ASAM OSI GroundTruth"

    _osi: betterosi.GroundTruth

    @classmethod
    def create(cls, gt: betterosi.GroundTruth):
        if len(gt.lane_boundary) == 0:
            raise RuntimeError("Empty Map")
        return cls(
            _osi=gt,
            lane_boundaries={b.id.value: LaneBoundaryOsi.create(b) for b in gt.lane_boundary},
            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):
        self.setup_lanes_and_boundaries()

    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 if i in map_osi_id2idx]
            l.predecessor_ids = [map_osi_id2idx[i] for i in l.predecessor_ids if i in map_osi_id2idx]
            l._map = self
            l.set_boundaries()
            l.set_polygon()

    def _to_binary_json(self):
        d = json.loads(self._osi.to_json())
        if "movingObject" in d:
            del d["movingObject"]
        return {b"osi": json.dumps(d).encode()}

    @classmethod
    def _from_binary_json(cls, d, **kwargs):
        gt = betterosi.GroundTruth().from_json(d[b"osi"].decode())
        if len(gt.lane_boundary) > 0:
            return cls.create(gt)
        else:
            return None

MapOsiCenterline dataclass

Bases: Map

Map representation based on ASAM OSI GroundTruth defining only the centerlines of lanes and nothing else. Does not conform to the omega-prime specification for Map.

Source code in omega_prime/map.py

@dataclass(repr=False)
class MapOsiCenterline(Map):
    "Map representation based on ASAM OSI GroundTruth defining only the centerlines of lanes and nothing else. Does not conform to the omega-prime specification for Map."

    _osi: betterosi.GroundTruth
    lanes: dict[int, LaneOsiCenterline]

    @classmethod
    def create(cls, gt: betterosi.GroundTruth, split_lanes: bool = False, split_lanes_length: float = 10, **kwargs):
        if len(gt.lane) == 0:
            raise RuntimeError("No Map")
        c = cls(
            _osi=gt,
            lanes={l.idx: l for l in [LaneOsiCenterline.create(l) for l in gt.lane]},
            lane_boundaries={},
        )
        if split_lanes:
            c._split(split_lanes_length)
        return c

    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
        self.align_predecessor_and_successor_relations()

    def _split(self, max_len: float):
        """
        Split lanes into segments of maximum length.

        This method post-processes the map by splitting lane centerlines that exceed
        the specified maximum length into smaller segments. It updates lane connections
        accordingly and removes connections between segments that are too far apart.

        Args:
            max_len (float): Maximum length allowed for each lane segment.
        """
        warn("The Postprocessing is ACTIVE! The lanes will be split into segments!!!")
        lanes_or = self.lanes
        lanes_new = {}
        idx_count = 0

        for lane in tqdm(lanes_or.values()):
            if lane.centerline.length > max_len:
                # Split the lane's centerline into segments of maximum length
                segments = split_linestring(lane.centerline, max_len)
            else:
                segments = [lane.centerline]

            # Create new lane objects for each segment
            segment_lanes = []
            for i, segment in enumerate(segments):
                # Create a copy of the lane with modified centerline
                # new_lane = copy.deepcopy(lane)

                new_lane = LaneOsiCenterline(
                    _osi=lane._osi,
                    idx=OsiLaneId(road_id=idx_count, lane_id=idx_count),
                    centerline=segment,
                    type=lane.type,
                    subtype=lane.subtype,
                    successor_ids=[],
                    predecessor_ids=[],
                )

                segment_lanes.append(new_lane)
                lanes_new[new_lane.idx.lane_id] = new_lane
                idx_count += 1

            for i, new_lane in enumerate(segment_lanes):
                if len(segments) == 1:
                    # If only one segment, keep original predecessors and successors
                    new_lane.predecessor_ids = lane.predecessor_ids
                    new_lane.successor_ids = lane.successor_ids
                elif i == 0:
                    # First segment: keep original predecessors, connect to next segment
                    new_lane.predecessor_ids = lane.predecessor_ids
                    new_lane.successor_ids = [segment_lanes[i + 1].idx]
                elif i == len(segments) - 1:
                    # Last segment: connect to previous segment, keep original successors
                    new_lane.predecessor_ids = [segment_lanes[i - 1].idx]
                    new_lane.successor_ids = lane.successor_ids
                else:
                    # Middle segments: connect to both neighbors
                    new_lane.predecessor_ids = [segment_lanes[i - 1].idx]
                    new_lane.successor_ids = [segment_lanes[i + 1].idx]

            # 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
                    idx = other_lane.successor_ids.index(lane.idx)
                    other_lane.successor_ids[idx] = segment_lanes[0].idx
                if lane.idx in other_lane.predecessor_ids:
                    # Replace reference to original lane with last segment
                    idx = other_lane.predecessor_ids.index(lane.idx)
                    other_lane.predecessor_ids[idx] = segment_lanes[-1].idx

        # Replace original lanes with segmented lanes

        # Do a check for the predecessor 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:
                        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:
                        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

        self.lanes = {lane.idx: lane for lane in lanes_new.values()}
        for lane in self.lanes.values():
            lane._map = self
        return self

    def _to_binary_json(self):
        d = json.loads(self._osi.to_json())
        if "movingObject" in d:
            del d["movingObject"]
        return {b"osi": json.dumps(d).encode()}

    @classmethod
    def _from_binary_json(cls, d, **kwargs):
        gt = betterosi.GroundTruth().from_json(d[b"osi"].decode())
        return cls.create(gt)

split_linestring(line, max_length)

Split a LineString into segments of maximum length.

Parameters:

Name	Type	Description	Default
line		shapely LineString to split	required
max_length		Maximum length of each segment	required

Returns:

Type	Description
	List of LineString segments

Source code in omega_prime/map.py

def split_linestring(line, max_length):
    """
    Split a LineString into segments of maximum length.

    Args:
        line: shapely LineString to split
        max_length: Maximum length of each segment

    Returns:
        List of LineString segments
    """
    segments = []

    # If line is already short enough, return it as is
    if line.length <= max_length:
        return [line]

    # Number of segments needed
    n_segments = int(np.ceil(line.length / max_length))

    # Get evenly spaced points along the line
    points = [line.interpolate(i / n_segments, normalized=True) for i in range(n_segments + 1)]

    # Create line segments
    for i in range(n_segments):
        segment_coords = [points[i].coords[0], points[i + 1].coords[0]]
        segments.append(shapely.LineString(segment_coords))

    return segments

MapOdr dataclass

Bases: Map

Source code in omega_prime/map_odr.py

@dataclass(repr=False)
class MapOdr(Map):
    odr_xml: str
    name: str
    step_size: float = 0.01
    _xodr_map: PyxodrRoadNetwork | None = None
    proj_string: str | None = None
    proj_offset: ProjectionOffset | None = None
    projection: pyproj.CRS | None = None
    _supported_file_suffixes = [".xodr", ".mcap", ".odr"]
    _binary_json_identifier = b"xodr"

    @property
    def xodr_map(self):
        if self._xodr_map is None:
            self.parse()
        return self._xodr_map

    @classmethod
    def from_file(
        cls,
        filename,
        topics: list[str] = ["/ground_truth_map", "ground_truth_map"],
        parse_map: bool = False,
        is_odr_xml: bool = False,
        is_mcap: bool = False,
        step_size=0.01,
        ignored_lane_types: set[str] = set([]),
        **kwargs,
    ):
        if Path(filename).suffix in [".xodr", ".odr"] or is_odr_xml:
            with open(filename) as f:
                odr_xml = f.read()
            return cls.create(
                odr_xml=odr_xml,
                name=Path(filename).stem,
                step_size=step_size,
                parse_map=parse_map,
                ignored_lane_types=ignored_lane_types,
            )
        if Path(filename).suffix in [".mcap"] or is_mcap:
            map = next(iter(betterosi.read(filename, mcap_topics=topics, mcap_return_betterosi=False)))
            return cls.create(
                odr_xml=map.open_drive_xml_content, name=map.map_reference, step_size=step_size, parse_map=parse_map
            )

    @property
    def lanes(self):
        if self._lanes is None:
            self.parse()
        return self._lanes

    @lanes.setter
    def lanes(self, val):
        self._lanes = val

    @property
    def lane_boundaries(self):
        if self._lane_boundaries is None:
            self.parse()
        return self._lane_boundaries

    @lane_boundaries.setter
    def lane_boundaries(self, val):
        self._lane_boundaries = val

    @classmethod
    def create(cls, odr_xml, name, step_size=0.01, parse_map: bool = False, ignored_lane_types: set[str] = set([])):
        self = cls(odr_xml=odr_xml, name=name, step_size=step_size, lanes={}, lane_boundaries={})
        self._lane_boundaries = None
        self._lanes = None
        self.ignored_lane_types = ignored_lane_types
        if parse_map:
            self.parse()
        return self

    def parse(self):
        rn = RoadNetwork(self.odr_xml, resolution=self.step_size, ignored_lane_types=self.ignored_lane_types)

        lane_boundaries = {}
        lanes = {}

        # Extract projection information from XML tree
        proj_string = None
        proj_offset = None
        projection = None

        # Get the header element from XML
        header = rn.tree.find("header")
        if header is not None:
            # Get geoReference if it exists
            geo_ref = header.find("geoReference")
            if geo_ref is not None and geo_ref.text:
                proj_string = geo_ref.text.strip()
                try:
                    projection = pyproj.CRS.from_proj4(proj_string)
                except pyproj.exceptions.CRSError as e:
                    logger.warning(f"Failed to parse projection string: {e}")

            # Get offset if it exists
            offset = header.find("offset")
            if offset is not None:
                try:
                    proj_offset = ProjectionOffset(
                        x=float(offset.get("x", "0")),
                        y=float(offset.get("y", "0")),
                        z=float(offset.get("z", "0")),
                        yaw=float(offset.get("hdg", "0")),
                    )
                except (ValueError, TypeError) as e:
                    logger.warning(f"Failed to parse offset: {e}")

        for road in rn.get_roads():
            lane_idx = 0
            for lane_section_id, lane_section in enumerate(road.lane_sections):
                for lane in lane_section.lanes:
                    boundary_line = getattr(lane, "boundary_line", None)
                    if boundary_line is None or not len(boundary_line):
                        logger.warning(
                            f"Skipping road {road.id} / lane_section {lane_section_id} / lane {lane.id}: missing boundary_line"
                        )
                        continue

                    try:
                        left_boundary = LaneBoundaryXodr.create(
                            lane, road.id, lane.id, lane_section_id, "left", lane_idx=lane_idx
                        )
                        right_boundary = LaneBoundaryXodr.create(
                            lane, road.id, lane.id, lane_section_id, "right", lane_idx=lane_idx
                        )
                    except Exception as e:
                        logger.error(
                            f"Failed to create boundaries for road {road.id} / lane_section {lane_section_id} / lane {lane.id}: {e}"
                        )
                        continue

                    lane_boundaries[left_boundary.idx] = left_boundary
                    lane_boundaries[right_boundary.idx] = right_boundary

                    try:
                        lane_obj = LaneXodr.create(lane, road, lane_section_id, lane_idx)
                        lanes[lane_obj.idx] = lane_obj
                    except Exception as e:
                        logger.error(
                            f"Failed to create lane object for road {road.id} / lane_section {lane_section_id} / lane {lane.id}: {e}"
                        )

                    lane_idx += 1

        self._xodr_map = rn
        self.lane_boundaries = lane_boundaries
        self.lanes = lanes
        self.proj_string = proj_string
        self.proj_offset = proj_offset
        self.projection = projection
        for lane in self.lanes.values():
            lane._map = self
            lane._set_boundaries()
            lane._set_polygon()
        for b in self._lane_boundaries.values():
            b._map = self

        return self

    def setup_lanes_and_boundaries(self):
        pass

    def to_file(self, filename: str | Path):
        """Export the current MapOdr to a .xodr file."""
        if isinstance(filename, str):
            filename = Path(filename)

        if filename.is_dir():
            filename = filename / f"{self.name}.xodr"

        if filename.suffix == "":
            filename = filename.with_suffix(".xodr")

        with open(filename, "w") as f:
            f.write(self.odr_xml)

    def to_osi(self):
        return MapAsamOpenDrive(map_reference=self.name, open_drive_xml_content=self.odr_xml)

    def _to_binary_json(self):
        return {b"xodr": self.odr_xml.encode(), b"xodr_name": self.name.encode()}

    @classmethod
    def _from_binary_json(cls, d, parse_map: bool = False, step_size: float = 0.01):
        return cls.create(
            odr_xml=d[b"xodr"].decode(),
            name=d[b"xodr_name"].decode(),
            parse_map=parse_map,
            step_size=step_size,
        )

to_file(filename)

Export the current MapOdr to a .xodr file.

Source code in omega_prime/map_odr.py

def to_file(self, filename: str | Path):
    """Export the current MapOdr to a .xodr file."""
    if isinstance(filename, str):
        filename = Path(filename)

    if filename.is_dir():
        filename = filename / f"{self.name}.xodr"

    if filename.suffix == "":
        filename = filename.with_suffix(".xodr")

    with open(filename, "w") as f:
        f.write(self.odr_xml)

Locator dataclass

Source code in omega_prime/locator.py

@dataclass(repr=False)
class Locator:
    all_lanes: Any  # array of all lanes
    external2internal_laneid: dict[Any, int] = field(init=False)
    internal2external_laneid: list[Any] = field(init=False)
    lane_point_distances: list = field(init=False)
    str_tree: shapely.STRtree = field(init=False)
    extended_centerlines: list[shapely.LineString] = field(init=False)

    g: nx.DiGraph = field(init=False)  # Lane Relation Graph

    @classmethod
    def from_map(cls, map):
        all_lanes = list(map.lanes.values())
        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]

        self.extended_centerlines = [ShapelyTrajectoryTools.extend_linestring(l.centerline) for l in self.all_lanes]
        if hasattr(self.all_lanes[0], "polygon") and self.all_lanes[0].polygon is not None:
            self.str_tree = shapely.STRtree([l.polygon for l in self.all_lanes])
        else:
            self.str_tree = shapely.STRtree([l.centerline for l in self.all_lanes])
        self.lane_point_distances = [
            np.unique(shapely.line_locate_point(cl, shapely.points(cl.coords))) for cl in self.extended_centerlines
        ]
        self.g = self._get_routing_graph()

    def get_route(self, start_id, end_id):
        return nx.shortest_path(self.g, start_id, end_id)

    def sts2xys(self, sts):
        xys = np.zeros((len(sts.s), 2), dtype=float) * np.nan
        l_ids = np.array([self.external2internal_laneid[i] for i in sts.roadlane_id.values])
        for l_id in set(l_ids):
            point_idxs = np.argwhere(l_ids == l_id)[:, 0]
            rel_sts = sts.isel(dict(time=point_idxs))
            l = self.extended_centerlines[l_id]
            xys[point_idxs, 0], xys[point_idxs, 1] = ShapelyTrajectoryTools.st2xy(
                l, rel_sts.s.values + ShapelyTrajectoryTools.l_append, rel_sts.t.values
            )
        return xys

    def xys2sts(self, xys, polygons=None):
        if isinstance(xys, np.ndarray) and xys.ndim == 2:
            assert xys.shape[1] == 2
            xys = shapely.points(xys)
        lat_distances, lon_distances = self._xys2sts(xys, polygons)
        single_lane_association = self.get_single_lane_association(lat_distances)
        sla = np.zeros(len(single_lane_association), dtype=tuple)
        for i, v in enumerate(single_lane_association):
            sla[i] = v
        sts = xr.Dataset(
            {
                "s": ("time", [lon_distances[lidx][i] for i, lidx in enumerate(single_lane_association)]),
                "t": ("time", [lat_distances[lidx][i] for i, lidx in enumerate(single_lane_association)]),
                "roadlane_id": ("time", sla),
            }
        )
        return sts

    def xys2lane_sts(self, lane_id, xys, internal_id=False):
        # xys should be an array of shapely objects or an array of points with dim (n_points, 2)
        # return (n_points, 2) where ret[:,0] is s and ret[:,1] is t
        if isinstance(xys, np.ndarray) and xys.ndim == 2:
            assert xys.shape[1] == 2
            xys = shapely.points(xys)
        lid = self.external2internal_laneid[lane_id] if not internal_id else lane_id
        lane_point_distances = self.lane_point_distances[lid]
        sts = ShapelyTrajectoryTools.xy2st(
            self.extended_centerlines[lid], x_or_xy=xys, line_point_distances=lane_point_distances
        )
        sts[:, 0] -= ShapelyTrajectoryTools.l_append
        return sts

    def _xys2sts(self, xys, polygons=None):
        # xys should be an array of shapely objects or an array of points with dim (n_points, 2)
        if isinstance(xys, np.ndarray) and xys.ndim == 2:
            assert xys.shape[1] == 2
            xys = shapely.points(xys)
        else:
            xys = np.array(xys)
        if polygons is None:
            polygons = xys
        lon_distances = defaultdict(lambda: np.nan * np.ones((len(xys),)))
        lat_distances = defaultdict(lambda: np.nan * np.ones((len(xys),)))
        point_idxs, intersection_lane_ids = self.str_tree.query(polygons, predicate="intersects")
        for l_id in set(intersection_lane_ids):
            lps = point_idxs[intersection_lane_ids == l_id]
            (
                lon_distances[self.internal2external_laneid[l_id]][lps],
                lat_distances[self.internal2external_laneid[l_id]][lps],
            ) = self.xys2lane_sts(l_id, xys[lps], internal_id=True).T
        try:
            no_associations = np.where(np.all(np.isnan(np.stack(list(lon_distances.values()))), axis=0))[0]
        except ValueError:
            # no arrays to stack
            no_associations = np.arange(len(xys))
        if hasattr(self.all_lanes[0], "polygon") and self.all_lanes[0].polygon is not None:
            no_asscociation_idxs, intersection_lane_ids = self.str_tree.query_nearest(polygons[no_associations])
        else:
            # Create an empty numpy array for no_asscociation_idxs
            no_asscociation_idxs = np.array([])
            intersection_lane_ids = np.array([])
            if len(no_associations) > 0:
                for idx, poly in enumerate(polygons[no_associations]):
                    # 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, [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
        no_asscociation_idxs = no_asscociation_idxs.astype(int)
        intersection_lane_ids = intersection_lane_ids.astype(int)
        for l_id in set(intersection_lane_ids):
            lps = no_associations[no_asscociation_idxs[intersection_lane_ids == l_id]]
            (
                lon_distances[self.internal2external_laneid[l_id]][lps],
                lat_distances[self.internal2external_laneid[l_id]][lps],
            ) = self.xys2lane_sts(l_id, xys[lps], internal_id=True).T

        no_asscociation_new = np.where(np.all(np.isnan(np.stack(list(lon_distances.values()))), axis=0))[0]

        assert len(no_asscociation_new) == 0
        return lat_distances, lon_distances

    def locate_mv(self, mv, use_polygon: bool = False):
        mv.polygon
        moving = mv._df.filter(pl.any_horizontal((pl.col("x", "y").diff() != 0).fill_null(True)).alias("is_moving"))[
            "total_nanos", "x", "y", "polygon"
        ]
        xrd = (
            self.xys2sts(moving["x", "y"].to_numpy(), polygons=moving["polygon"] if use_polygon else None)
            .assign_coords({"time": moving["total_nanos"].to_numpy()})
            .set_coords("time")
        )
        if moving.height < mv._df.height:
            xrd = xrd.sel({"time": mv._df["total_nanos"].to_numpy()}, method="ffill", drop=True)
            xrd["time"] = mv._df["total_nanos"].to_numpy()
        return xrd

    def query_centerlines(self, point, range_percentage=0.1):
        """
        Query the nearest centerline and all centerlines within a range percentage.

        :param point: A shapely Point object representing the query location.
        :param range_percentage: The range as a percentage of the total length of the nearest centerline. Default is 0.1 (10%).
        :return: A NDArray with all the Lane Idx in the Range.
        """
        # Query the nearest centerline
        nearest_idx = self.str_tree.query_nearest(point)
        nearest_centerline = self.extended_centerlines[nearest_idx[0]]

        # Calculate the range based on the nearest centerline's length
        range_distance = nearest_centerline.distance(point) * (1 + range_percentage)

        # Create a buffer around the point
        buffer = point.buffer(range_distance)

        # Query all centerlines within the buffer
        nearby_idxs = self.str_tree.query(buffer, predicate="intersects")

        # If there was no intersection, return the nearest centerline
        if nearby_idxs.size == 0:
            return nearest_idx

        return nearby_idxs

    def _get_routing_graph(self):
        all_lanes = self.all_lanes
        str_tree = self.str_tree
        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)
                except KeyError:
                    pass
            for external_sid in lane.successor_ids:
                try:
                    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
            ]
            left_neigbours = [
                int(i) for i in str_tree.query(lane.left_boundary.polyline, predicate="covered_by") if int(i) != lid
            ]
            for rn in right_neigbours:
                g.add_edge(lid, rn, label=LaneRelation.neighbour_right)
            for ln in left_neigbours:
                g.add_edge(lid, ln, label=LaneRelation.neighbour_left)
        return g

    def get_single_lane_association(
        self, traveler_lane_intersections: dict[Any, Any], overlaps: None | dict[Any, float] = None
    ):
        """
        filter traveling path of traveler, so that traveler is not assigned to lanes that are only reachable through a merging or crossing relation
        return format: road, lane
        """
        import networkx as nx

        g = nx.Graph()
        nodes = defaultdict(list)
        for external_lid, v in traveler_lane_intersections.items():
            for timeidx in np.where(~np.isnan(v))[0]:
                nodes[timeidx].append(self.external2internal_laneid[external_lid])
        nodes_per_time = [nodes[i] for i in range(len(nodes))]
        g.add_node("start", pos=(-1, -1))
        g.add_node("end", pos=(len(nodes_per_time), -1))
        for i, nodes_of_time in enumerate(nodes_per_time[:-1]):
            for n in nodes_of_time:
                for next_n in nodes_per_time[i + 1]:
                    g.add_node((n, i), pos=(i, n))
                    g.add_node((next_n, i + 1), pos=(i + 1, next_n))
                    try:
                        if n == next_n:
                            if overlaps is None:
                                weight = 1
                            else:
                                weight = 1 - overlaps[self.internal2external_laneid[n]][i + 1]
                        elif any(
                            [
                                LaneRelation.neighbour_left in o or LaneRelation.neighbour_right in o
                                for o in self.g.get_edge_data(n, next_n)["label"]
                                + self.g.get_edge_data(next_n, n)["label"]
                            ]
                        ):
                            weight = 2
                        elif self.g.get_edge_data(n, next_n)["label"] in [
                            LaneRelation.predecessor,
                            LaneRelation.successor,
                        ] or self.g.get_edge_data(n, next_n)["label"] in [
                            LaneRelation.predecessor,
                            LaneRelation.successor,
                        ]:
                            weight = 2
                        else:
                            weight = 3
                    except Exception:
                        weight = 4
                    g.add_edge((n, i), (next_n, i + 1), weight=weight)
        for n in nodes_per_time[0]:
            g.add_edge("start", (n, 0), weight=1)
        for n in nodes_per_time[-1]:
            g.add_edge((n, len(nodes_per_time) - 1), "end", weight=1)
        sp = nx.shortest_path(g, "start", "end", weight="weight")[1:-1]
        fixed_traveler_path = [self.internal2external_laneid[o[0]] for o in sp]

        overlaps = [traveler_lane_intersections[lid][i] for i, lid in enumerate(fixed_traveler_path)]
        assert not np.any(np.isnan(overlaps))
        return fixed_traveler_path

    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]

get_single_lane_association(traveler_lane_intersections, overlaps=None)

filter traveling path of traveler, so that traveler is not assigned to lanes that are only reachable through a merging or crossing relation return format: road, lane

Source code in omega_prime/locator.py

def get_single_lane_association(
    self, traveler_lane_intersections: dict[Any, Any], overlaps: None | dict[Any, float] = None
):
    """
    filter traveling path of traveler, so that traveler is not assigned to lanes that are only reachable through a merging or crossing relation
    return format: road, lane
    """
    import networkx as nx

    g = nx.Graph()
    nodes = defaultdict(list)
    for external_lid, v in traveler_lane_intersections.items():
        for timeidx in np.where(~np.isnan(v))[0]:
            nodes[timeidx].append(self.external2internal_laneid[external_lid])
    nodes_per_time = [nodes[i] for i in range(len(nodes))]
    g.add_node("start", pos=(-1, -1))
    g.add_node("end", pos=(len(nodes_per_time), -1))
    for i, nodes_of_time in enumerate(nodes_per_time[:-1]):
        for n in nodes_of_time:
            for next_n in nodes_per_time[i + 1]:
                g.add_node((n, i), pos=(i, n))
                g.add_node((next_n, i + 1), pos=(i + 1, next_n))
                try:
                    if n == next_n:
                        if overlaps is None:
                            weight = 1
                        else:
                            weight = 1 - overlaps[self.internal2external_laneid[n]][i + 1]
                    elif any(
                        [
                            LaneRelation.neighbour_left in o or LaneRelation.neighbour_right in o
                            for o in self.g.get_edge_data(n, next_n)["label"]
                            + self.g.get_edge_data(next_n, n)["label"]
                        ]
                    ):
                        weight = 2
                    elif self.g.get_edge_data(n, next_n)["label"] in [
                        LaneRelation.predecessor,
                        LaneRelation.successor,
                    ] or self.g.get_edge_data(n, next_n)["label"] in [
                        LaneRelation.predecessor,
                        LaneRelation.successor,
                    ]:
                        weight = 2
                    else:
                        weight = 3
                except Exception:
                    weight = 4
                g.add_edge((n, i), (next_n, i + 1), weight=weight)
    for n in nodes_per_time[0]:
        g.add_edge("start", (n, 0), weight=1)
    for n in nodes_per_time[-1]:
        g.add_edge((n, len(nodes_per_time) - 1), "end", weight=1)
    sp = nx.shortest_path(g, "start", "end", weight="weight")[1:-1]
    fixed_traveler_path = [self.internal2external_laneid[o[0]] for o in sp]

    overlaps = [traveler_lane_intersections[lid][i] for i, lid in enumerate(fixed_traveler_path)]
    assert not np.any(np.isnan(overlaps))
    return fixed_traveler_path

query_centerlines(point, range_percentage=0.1)

Query the nearest centerline and all centerlines within a range percentage.

:param point: A shapely Point object representing the query location. :param range_percentage: The range as a percentage of the total length of the nearest centerline. Default is 0.1 (10%). :return: A NDArray with all the Lane Idx in the Range.

Source code in omega_prime/locator.py

def query_centerlines(self, point, range_percentage=0.1):
    """
    Query the nearest centerline and all centerlines within a range percentage.

    :param point: A shapely Point object representing the query location.
    :param range_percentage: The range as a percentage of the total length of the nearest centerline. Default is 0.1 (10%).
    :return: A NDArray with all the Lane Idx in the Range.
    """
    # Query the nearest centerline
    nearest_idx = self.str_tree.query_nearest(point)
    nearest_centerline = self.extended_centerlines[nearest_idx[0]]

    # Calculate the range based on the nearest centerline's length
    range_distance = nearest_centerline.distance(point) * (1 + range_percentage)

    # Create a buffer around the point
    buffer = point.buffer(range_distance)

    # Query all centerlines within the buffer
    nearby_idxs = self.str_tree.query(buffer, predicate="intersects")

    # If there was no intersection, return the nearest centerline
    if nearby_idxs.size == 0:
        return nearest_idx

    return nearby_idxs

get_lane_centerline(right_border, left_border)

middle line between (interpolated) boundaries, oriented in direction of lane

Source code in omega_prime/locator.py

def get_lane_centerline(right_border: shapely.LineString, left_border: shapely.LineString) -> shapely.LineString:
    """middle line between (interpolated) boundaries, oriented in direction of lane"""
    ses = np.unique(
        np.concatenate(
            [
                shapely.line_locate_point(left_border, shapely.points(right_border.coords), normalized=True),
                shapely.line_locate_point(right_border, shapely.points(left_border.coords), normalized=True),
            ]
        )
    )

    points = np.zeros((len(ses), 2))
    for i, (rbp, lbp) in enumerate(
        zip(right_border.interpolate(ses, normalized=True), left_border.interpolate(ses, normalized=True))
    ):
        points[i, :] = shapely.MultiPoint([rbp, lbp]).minimum_rotated_rectangle.centroid.coords

    # TODO some smoothing operation could be helpful
    cl = shapely.LineString(points)

    if cl.is_empty or not cl.is_valid:
        raise RuntimeError("Could not compute centerline for lane!")
    return cl

DatasetConverter

Bases: ABC

Source code in omega_prime/converters/converter.py

class DatasetConverter(ABC):
    def __init__(self, dataset_path: str, out_path: str = "./", n_workers=1) -> None:
        self._dataset_path = Path(dataset_path)
        self._out_path = Path(out_path)
        self.n_workers = n_workers
        self.len = None

    @abstractmethod
    def get_source_recordings(self) -> list:
        """
        Abstract method to get a list of the source recordings.
        The method should be implemented in subclasses to handle specific dataset formats.
        Returns:
            source_recordings: List of the source recordings. Could be of any type as further processed in get_recordings.
        """
        pass

    @abstractmethod
    def get_recordings(self, source_recording) -> Iterator:
        """
        Abstract method to get all recordings in a source-recording-instance of the specific dataset.
        The method should be implemented in subclasses to handle specific dataset formats.
        Args:
            source_recordings: List of the source recordings. Could be of any type as returned by get_source_recordings.
        Yields:
            recording: Each recording in the source-recording-instance, one at a time. Could be of any type as further processed in to_omega_prime_recording and get_recording_id.
        """
        pass

    @abstractmethod
    def to_omega_prime_recording(self, recording) -> Recording:
        """
        Abstract method to convert a raw recording into an omega prime recording instance.
        The method should be implemented in subclasses to handle specific dataset formats.
        Args:
            recording: A recording of any type as returned by get_omega_prime_recordings.
        Returns:
            Recording: An instance of the Recording class containing the processed data.
        """
        pass

    @abstractmethod
    def get_recording_name(self, recording) -> str:
        """
        Abstract method to get the name for a given recording.
        The method should be implemented in subclasses to handle specific dataset formats.
        Args:
            recording: Recording of any type as returned by get_recordings.
        Returns:
            str: unique name of recording.
        """
        pass

    def convert_source_recording(
        self, source_recording, save_as_parquet: bool = False, skip_existing: bool = False, log_file: Path | None = None
    ) -> None:
        try:
            for recording in self.get_recordings(source_recording):
                out_filename = (
                    self._out_path / f"{self.get_recording_name(recording)}.{'parquet' if save_as_parquet else 'mcap'}"
                )
                status = Status(str(source_recording), str(out_filename))
                if not skip_existing or not out_filename.exists():
                    Path(out_filename).parent.mkdir(exist_ok=True, parents=True)
                    try:
                        rec = self.to_omega_prime_recording(recording)
                        status.set_success()
                    except Exception as e:
                        logger.error(
                            f"Error converting recording {self.get_recording_name(recording)}: {traceback.format_exc()}"
                        )
                        rec = None
                        status.set_error(str(e))
                    else:
                        try:
                            if save_as_parquet:
                                rec.to_parquet(out_filename)
                            else:
                                rec.to_mcap(out_filename)
                        except Exception as e:
                            logger.error(
                                f"Error saving recording {self.get_recording_name(recording)}: {traceback.format_exc()}"
                            )
                            status.set_error(e)
                else:
                    status.set_skip()

                if log_file is not None:
                    with FileLock(log_file.with_suffix(".csv.lock")):
                        status.write(log_file)

        except Exception as e:
            logger.error(f"Error processing source recording {source_recording}: {e} - {traceback.format_exc()}")
            raise e

    def convert(
        self,
        n_workers: int | None = None,
        save_as_parquet: bool = False,
        skip_existing: bool = False,
        write_log: bool = False,
    ) -> None:
        if n_workers is None:
            n_workers = self.n_workers
        if n_workers == -1:
            n_workers = jb.cpu_count() - 1
        self._out_path.mkdir(exist_ok=True, parents=True)
        recordings = self.get_source_recordings()

        # Create a log file if requested
        log_file = None
        if write_log:
            log_file = self._out_path / "conversion_log.csv"
            with open(log_file, "w", newline="") as csvfile:
                fieldnames = ["file_path_input", "status", "file_path_output", "error_message"]
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
        if self.len is None:
            try:
                self.len = len(recordings)
            except TypeError:
                pass
        if n_workers > 1:
            partial_fct = partial(
                self.convert_source_recording,
                save_as_parquet=save_as_parquet,
                skip_existing=skip_existing,
                log_file=log_file,
            )
            with tqdm_joblib(desc="Source Recordings", total=self.len):
                jb.Parallel(n_jobs=n_workers)(jb.delayed(partial_fct)(rec) for rec in recordings)
        else:
            for rec in tqdm(recordings, total=self.len):
                self.convert_source_recording(
                    rec, save_as_parquet=save_as_parquet, skip_existing=skip_existing, log_file=log_file
                )

    def yield_recordings(self) -> Iterator[Recording]:
        source_recordings = self.get_source_recordings()
        for sr in tqdm(source_recordings, total=len(source_recordings)):
            for recording in self.get_recordings(sr):
                yield self.to_omega_prime_recording(recording)

    @classmethod
    def convert_cli(
        cls,
        dataset_path: Annotated[
            Path,
            typer.Argument(exists=True, dir_okay=True, file_okay=True, readable=True, help="Root of the dataset"),
        ],
        output_path: Annotated[
            Path,
            typer.Argument(
                file_okay=False, writable=True, help="In which folder to write the created omega-prime files"
            ),
        ],
        n_workers: Annotated[int, typer.Option(help="Set to -1 for n_cpus-1 workers.")] = 1,
        save_as_parquet: Annotated[
            bool,
            typer.Option(
                help="If activated, omega-prime recordings will be stored as parquet files instead of mcap (use for large recordings). Will loose information in OSI that are not mandatory in omega-prime."
            ),
        ] = False,
        skip_existing: Annotated[bool, typer.Option(help="Only convert not yet converted files")] = False,
        write_log: Annotated[bool, typer.Option(help="Write a log file with the conversion process")] = False,
    ):
        Path(output_path).mkdir(exist_ok=True)
        cls(dataset_path=dataset_path, out_path=output_path, n_workers=n_workers).convert(
            save_as_parquet=save_as_parquet, skip_existing=skip_existing, write_log=write_log
        )

get_recording_name(recording) abstractmethod

Abstract method to get the name for a given recording. The method should be implemented in subclasses to handle specific dataset formats. Args: recording: Recording of any type as returned by get_recordings. Returns: str: unique name of recording.

Source code in omega_prime/converters/converter.py

@abstractmethod
def get_recording_name(self, recording) -> str:
    """
    Abstract method to get the name for a given recording.
    The method should be implemented in subclasses to handle specific dataset formats.
    Args:
        recording: Recording of any type as returned by get_recordings.
    Returns:
        str: unique name of recording.
    """
    pass

get_recordings(source_recording) abstractmethod

Abstract method to get all recordings in a source-recording-instance of the specific dataset. The method should be implemented in subclasses to handle specific dataset formats. Args: source_recordings: List of the source recordings. Could be of any type as returned by get_source_recordings. Yields: recording: Each recording in the source-recording-instance, one at a time. Could be of any type as further processed in to_omega_prime_recording and get_recording_id.

Source code in omega_prime/converters/converter.py

@abstractmethod
def get_recordings(self, source_recording) -> Iterator:
    """
    Abstract method to get all recordings in a source-recording-instance of the specific dataset.
    The method should be implemented in subclasses to handle specific dataset formats.
    Args:
        source_recordings: List of the source recordings. Could be of any type as returned by get_source_recordings.
    Yields:
        recording: Each recording in the source-recording-instance, one at a time. Could be of any type as further processed in to_omega_prime_recording and get_recording_id.
    """
    pass

get_source_recordings() abstractmethod

Abstract method to get a list of the source recordings. The method should be implemented in subclasses to handle specific dataset formats. Returns: source_recordings: List of the source recordings. Could be of any type as further processed in get_recordings.

Source code in omega_prime/converters/converter.py

@abstractmethod
def get_source_recordings(self) -> list:
    """
    Abstract method to get a list of the source recordings.
    The method should be implemented in subclasses to handle specific dataset formats.
    Returns:
        source_recordings: List of the source recordings. Could be of any type as further processed in get_recordings.
    """
    pass

to_omega_prime_recording(recording) abstractmethod

Abstract method to convert a raw recording into an omega prime recording instance. The method should be implemented in subclasses to handle specific dataset formats. Args: recording: A recording of any type as returned by get_omega_prime_recordings. Returns: Recording: An instance of the Recording class containing the processed data.

Source code in omega_prime/converters/converter.py

@abstractmethod
def to_omega_prime_recording(self, recording) -> Recording:
    """
    Abstract method to convert a raw recording into an omega prime recording instance.
    The method should be implemented in subclasses to handle specific dataset formats.
    Args:
        recording: A recording of any type as returned by get_omega_prime_recordings.
    Returns:
        Recording: An instance of the Recording class containing the processed data.
    """
    pass

Metric dataclass

Class to compute metrics based on polars dataframes.

Source code in omega_prime/metrics.py

@dataclass
class Metric:
    """Class to compute metrics based on polars dataframes."""

    compute_func: Callable[[pl.LazyFrame, ...], tuple[pl.LazyFrame, dict[str, pl.LazyFrame]]]
    """The function that actually computes the metric"""
    computes_columns: list[str] = field(default_factory=list)
    """Names of the columns that will be added to the dataframe by this metrics"""
    computes_properties: list[str] = field(default_factory=list)
    """Keys of the tables added to the properties dictionary by this metric"""
    requires_columns: list[str] = field(default_factory=list)
    """Columns that must be present in the dataframe before this metric can be calculated"""
    requires_properties: list[str] = field(default_factory=list)
    """Keys of tables that must be present in the properties dictionary before this metric can be calculated."""
    computes_intermediate_columns: list[str] = field(default_factory=list)
    """Same as computes_columns by these ones will not be returned in the end and are only available to other metrics"""
    computes_intermediate_properties: list[str] = field(default_factory=list)
    """ Same as computes_properties but these ones will not be returned in the end and are only available to other metrics."""
    _parameters: list = field(init=False)
    """All parameters of metrics that need to be set on computation"""

    def compute_lazy(self, df: pl.LazyFrame, **kwargs) -> tuple[pl.LazyFrame, dict[str, pl.LazyFrame]]:
        try:
            df, properties = self.compute_func(df, **kwargs)
            assert isinstance(df, pl.LazyFrame)
            assert all(p in properties for p in self.computes_properties + self.computes_intermediate_properties)
            return df, properties

        except TypeError as e:
            raise TypeError(
                f"Missing parameter for Metric with compute_func {self.compute_func.__name__}: {repr(e)}"
            ) from e

    def __post_init__(self):
        sig = inspect.signature(self.compute_func)
        parameters = sig.parameters
        assert "df" in parameters
        assert all(p in parameters for p in self.requires_properties)

        self._parameters = [
            v for k, v in parameters.items() if k not in ["df", "args", "kwargs"] + self.requires_properties
        ]

    def __call__(self, df: pl.DataFrame, **kwargs):
        try:
            if not isinstance(df, pl.LazyFrame):
                df = pl.LazyFrame(df)
            return self.compute_lazy(df, **kwargs)
        except TypeError as e:
            raise TypeError(
                f"Missing paramter for Metric with compute_func {self.compute_func.__name__}: {repr(e)}"
            ) from e

compute_func instance-attribute

The function that actually computes the metric

computes_columns = field(default_factory=list) class-attribute instance-attribute

Names of the columns that will be added to the dataframe by this metrics

computes_intermediate_columns = field(default_factory=list) class-attribute instance-attribute

Same as computes_columns by these ones will not be returned in the end and are only available to other metrics

computes_intermediate_properties = field(default_factory=list) class-attribute instance-attribute

Same as computes_properties but these ones will not be returned in the end and are only available to other metrics.

computes_properties = field(default_factory=list) class-attribute instance-attribute

Keys of the tables added to the properties dictionary by this metric

requires_columns = field(default_factory=list) class-attribute instance-attribute

Columns that must be present in the dataframe before this metric can be calculated

requires_properties = field(default_factory=list) class-attribute instance-attribute

Keys of tables that must be present in the properties dictionary before this metric can be calculated.

MetricManager dataclass

Source code in omega_prime/metrics.py

@dataclass
class MetricManager:
    metrics: list[Metric] = field(default_factory=lambda: metrics)
    """List of metrics to compute"""
    exclude_columns: list[str] = field(default_factory=list)
    """List of columns computed by the metrics that do not need to be computed"""
    exclude_properties: list[str] = field(default_factory=list)
    """List of tables in the properties dict that do not need to be computed"""
    _dependencies: dict[int | str, list[int | str]] = field(init=False)
    """Automatically derived dependencies between metrics"""
    _ordered_metrics: list[Metric] = field(init=False)
    """Automatically derived execution order of metrics"""
    _parameters: list = field(init=False)
    """Automatically derived list of parameters to keep"""

    def __post_init__(self):
        self._dependencies = {
            val: [i]
            for i, m in enumerate(self.metrics)
            for val in [f"column_{n}" for n in m.computes_columns + m.computes_intermediate_columns]
            + [f"property_{n}" for n in m.computes_properties + m.computes_intermediate_properties]
        } | {
            i: [f"column_{n}" for n in m.requires_columns] + [f"property_{n}" for n in m.requires_properties]
            for i, m in enumerate(self.metrics)
        }

        unresovled_dependencies = {
            k: v for k, vv in self._dependencies.items() for v in vv if v not in self._dependencies
        }
        if len(unresovled_dependencies) > 0:
            error_dict = {f"self.metrics[{k}]": v for k, v in unresovled_dependencies.items()}
            raise RuntimeError(
                f"There are columns and properties required by metrics, that are never computed: {error_dict}"
            )

        self._parameters = [v for m in self.metrics for v in m._parameters]

        self.exclude_columns += [v for m in self.metrics for v in m.computes_intermediate_columns]
        self.exclude_properties += [v for m in self.metrics for v in m.computes_intermediate_properties]

        ts = graphlib.TopologicalSorter(self._dependencies)
        self._ordered_metrics = [self.metrics[o] for o in ts.static_order() if isinstance(o, int)]

    def __repr__(self):
        return f"computes columns: {[c for m in self._ordered_metrics for c in m.computes_columns]} - computes properties {[p for m in self._ordered_metrics for p in m.computes_properties]} - parameters {list(set([str(m) for m in self._parameters]))}"

    def compute(self, r: Recording, **kwargs) -> tuple[pl.DataFrame, dict[str, pl.DataFrame]]:
        if "polygon" not in r._df.columns:
            r._df = r._add_polygons(r._df)
        if "geometry" not in r._df.columns:
            r._df = r._df.with_columns(geometry=st.from_shapely("polygon"))

        df = pl.LazyFrame(r._df)
        properties = {}
        for m in self._ordered_metrics:
            df, new_p = m.compute_lazy(
                df=df,
                **{k: properties[k] for k in m.requires_properties},
                **{k: v for k, v in kwargs.items() if k in [p.name for p in m._parameters]},
            )
            properties |= new_p
        for k in self.exclude_properties:
            del properties[k]
        df = df.drop(self.exclude_columns)
        res = pl.collect_all([df] + list(properties.values()))
        df, computed_props = res[0], res[1:]
        assert all(c in df.columns or c in self.exclude_columns for m in self.metrics for c in m.computes_columns)
        return df, {k: v for k, v in zip(properties.keys(), computed_props)}

    def plot_dependencies(self):
        import networkx as nx
        import matplotlib.pyplot as plt

        i = 0
        pos = {}
        G = nx.DiGraph()

        for m in self._ordered_metrics:
            n = m.compute_func.__name__
            pos[n] = [0, -i]
            i += 1
            cn = [f"column_{c}" for c in m.computes_columns + m.computes_intermediate_columns] + [
                f"property_{c}" for c in m.computes_properties + m.computes_intermediate_properties
            ]
            pos |= {k: [1 + j, -i] for j, k in enumerate(cn)}
            G.add_node(n, color="lightblue")
            for c in cn:
                G.add_node(c, color="lightgreen")
                G.add_edge(n, c, label="computes")
            for r in [f"column_{c}" for c in m.requires_columns] + [f"property_{p}" for p in m.requires_properties]:
                G.add_edge(r, n, label="required by")
            i += 1

        # Draw nodes and edges
        fig, ax = plt.subplots()
        nx.draw(
            G,
            pos,
            with_labels=True,
            node_size=2000,
            node_color=list(nx.get_node_attributes(G, "color").values()),
            arrows=True,
            font_size=8,
            ax=ax,
        )

        return fig

exclude_columns = field(default_factory=list) class-attribute instance-attribute

List of columns computed by the metrics that do not need to be computed

exclude_properties = field(default_factory=list) class-attribute instance-attribute

List of tables in the properties dict that do not need to be computed

metrics = field(default_factory=(lambda: metrics)) class-attribute instance-attribute

List of metrics to compute

distance_traveled(df)

Metric that computes the column distance_traveled

Source code in omega_prime/metrics.py

@metric(computes_columns=["distance_traveled"])
def distance_traveled(df) -> tuple[pl.DataFrame, dict[str, pl.DataFrame]]:
    """Metric that computes the column `distance_traveled`"""
    return df.with_columns(
        (pl.col("x").diff() ** 2 + pl.col("y").diff() ** 2)
        .sqrt()
        .fill_null(0.0)
        .cum_sum()
        .over("idx", order_by="total_nanos")
        .alias("distance_traveled"),
    ), {}

metric(computes_columns=None, computes_properties=None, requires_columns=None, requires_properties=None, computes_intermediate_columns=None, computes_intermediate_properties=None)

Decorator to turn a function into a Metric

Source code in omega_prime/metrics.py

def metric(
    computes_columns: list[str] | None = None,
    computes_properties: list[str] | None = None,
    requires_columns: list[str] | None = None,
    requires_properties: list[str] | None = None,
    computes_intermediate_columns: list[str] | None = None,
    computes_intermediate_properties: list[str] | None = None,
):
    """Decorator to turn a function into a Metric"""

    def decorator(func):
        return Metric(
            compute_func=func,
            computes_columns=computes_columns or [],
            computes_properties=computes_properties or [],
            requires_columns=requires_columns or [],
            requires_properties=requires_properties or [],
            computes_intermediate_columns=computes_intermediate_columns or [],
            computes_intermediate_properties=computes_intermediate_properties or [],
        )

    return decorator

p_timegaps_and_min_p_timgaps(df, /, ego_id, crossed, timegaps, time_buffer=2000000000.0)

Metrics that computes a predicted timegap between ego_id and all other objects. time_buffer gives the timespan in which intersection of trajectories is tested. The prediction is based on constant velocity following the same trajectory as observed.

Source code in omega_prime/metrics.py

@metric(
    requires_columns=["distance_traveled", "vel"],
    requires_properties=["crossed", "timegaps"],
    computes_properties=["p_timegaps", "min_p_timegaps"],
)
def p_timegaps_and_min_p_timgaps(df, /, ego_id, crossed, timegaps, time_buffer=2e9):
    """Metrics that computes a predicted timegap between `ego_id` and all other objects. `time_buffer` gives the timespan in which intersection of trajectories is tested. The prediction is based on constant velocity following the same trajectory as observed."""
    p_timegaps = (
        crossed.join(timegaps, how="right", suffix="_overlap", on=["idx", "idx_ego"])
        .with_columns(
            pl.when(pl.col("total_nanos") >= pl.col("total_nanos_overlap"))
            .then((pl.col("total_nanos_overlap") - pl.col("total_nanos")) / 1e9)
            .otherwise((pl.col("distance_traveled_overlap") - pl.col("distance_traveled")) / pl.col("vel"))
            .alias("time_to_overlap"),
            pl.when(pl.col("total_nanos_ego") >= pl.col("total_nanos_ego_overlap"))
            .then((pl.col("total_nanos_ego_overlap") - pl.col("total_nanos_ego")) / 1e9)
            .otherwise((pl.col("distance_traveled_ego_overlap") - pl.col("distance_traveled_ego")) / pl.col("vel_ego"))
            .alias("time_to_overlap_ego"),
        )
        .with_columns(
            -(
                pl.col("time_to_overlap_ego")
                - pl.col("time_to_overlap")
                + (pl.col("total_nanos_ego") - pl.col("total_nanos")) / 1e9
            ).alias("p_timegap")
        )
        .group_by("idx_ego", "idx", "total_nanos_ego")
        .agg(
            pl.col("p_timegap", "total_nanos")
            .sort_by(pl.col("p_timegap").abs(), descending=False, nulls_last=True)
            .first()
        )
        .sort("idx_ego", "idx", "total_nanos_ego")
    )

    min_p_timegaps = p_timegaps.group_by("idx_ego", "idx").agg(
        pl.col("p_timegap").sort_by(pl.col("p_timegap").abs(), descending=False).first()
    )

    return df, {
        "p_timegaps": p_timegaps,
        "min_p_timegaps": min_p_timegaps,
    }

timegaps_and_min_timgaps(df, /, ego_id, time_buffer=2000000000.0)

Metrics that computes timegaps between ego_id and all other objects. time_buffer gives the timespan in which intersection of trajectories is tested

Source code in omega_prime/metrics.py

@metric(
    requires_columns=["distance_traveled", "vel"],
    computes_properties=["timegaps", "min_timegaps"],
    computes_intermediate_properties=["crossed"],
)
def timegaps_and_min_timgaps(df, /, ego_id, time_buffer=2e9):
    """Metrics that computes timegaps between `ego_id` and all other objects. `time_buffer` gives the timespan in which intersection of trajectories is tested"""
    ego_df = df.filter(idx=ego_id)

    crossed = df.join(ego_df, how="cross", suffix="_ego")

    crossed = crossed.filter(
        (pl.col("total_nanos_ego") - time_buffer) <= pl.col("total_nanos"),
        (pl.col("total_nanos_ego") + time_buffer) >= pl.col("total_nanos"),
        pl.col("idx_ego") != pl.col("idx"),
    )

    all_timegaps = (
        crossed.filter(pl.col("geometry").st.intersects(pl.col("geometry_ego")))
        .with_columns(timegap=(pl.col("total_nanos") - pl.col("total_nanos_ego")) / 1e9)
        .select(
            "idx_ego", "idx", "total_nanos_ego", "total_nanos", "timegap", "distance_traveled", "distance_traveled_ego"
        )
    )

    timegaps = (
        all_timegaps.group_by("idx", "idx_ego", "total_nanos_ego")
        .agg(
            pl.col("timegap", "total_nanos", "distance_traveled", "distance_traveled_ego").get(
                pl.col("timegap").abs().arg_min()
            ),
        )
        .sort("idx_ego", "idx", "total_nanos_ego")
        .select(
            "idx_ego", "idx", "total_nanos_ego", "timegap", "total_nanos", "distance_traveled", "distance_traveled_ego"
        )
    )
    min_timegaps = timegaps.group_by("idx_ego", "idx").agg(
        pl.col("timegap").get(pl.col("timegap").abs().arg_min()).alias("min_timegap")
    )

    return df, {"timegaps": timegaps, "min_timegaps": min_timegaps, "crossed": crossed}

vel(df)

Metric that computes the column length of the speed vecotr vel

Source code in omega_prime/metrics.py

@metric(computes_columns=["vel"])
def vel(df) -> tuple[pl.DataFrame, dict[str, pl.DataFrame]]:
    """Metric that computes the column length of the speed vecotr `vel`"""
    return df.with_columns(
        (pl.col("vel_x") ** 2 + pl.col("vel_y") ** 2).sqrt().alias("vel"),
    ), {}

MapSegmentType

Bases: Enum

Classification of MapSegments.

Source code in omega_prime/mapsegment.py

class MapSegmentType(Enum):
    """Classification of MapSegments."""

    STRAIGHT = "straight"
    JUNCTION = "junction"
    ROUNDABOUT = "roundabout"
    RAMP_ON = "ramp_on"
    RAMP_OFF = "ramp_off"
    UNKNOWN = "unknown"

MapSegmentation

Bases: ABC

Abstract base class for map segmentation that handles multiple segments on a single map. Concrete implementations must define how to extract lane-specific information.

Source code in omega_prime/mapsegment.py

class MapSegmentation(ABC):
    """
    Abstract base class for map segmentation that handles multiple segments on a single map.
    Concrete implementations must define how to extract lane-specific information.
    """

    def __init__(self, recording, concave_hull_ratio=0.3):
        self.map = recording.map
        self.lanes = recording.map.lanes
        self.trafficlight = {}
        self.trafficlight_ids = set()
        self.intersections = []
        self.lane_dict = {}
        self.lane_successors_dict = {}
        self.lane_predecessors_dict = {}
        self.intersecting_lanes_dict = {}
        self.intersection_dict = {}
        self.lane_segment_dict = {}
        self.segments = []
        self.concave_hull_ratio = concave_hull_ratio

        segment_name = nt("SegmentName", ["lane_id", "segment_idx", "segment"])
        for lane in self.lanes.values():
            self.lane_segment_dict[self._get_lane_id(lane)] = segment_name(self._get_lane_id(lane), None, None)

    @abstractmethod
    def _get_lane_id(self, lane) -> Any:
        """Extract lane ID from a lane object. Map-type specific."""
        pass

    @abstractmethod
    def _get_lane_centerline(self, lane) -> shapely.LineString:
        """Extract centerline from a lane object. Map-type specific."""
        pass

    @abstractmethod
    def _get_lane_successors(self, lane) -> list:
        """Extract successor IDs from a lane object. Map-type specific."""
        pass

    @abstractmethod
    def _get_lane_predecessors(self, lane) -> list:
        """Extract predecessor IDs from a lane object. Map-type specific."""
        pass

    @abstractmethod
    def _has_traffic_light(self, lane) -> bool:
        """Check if lane has traffic light. Map-type specific."""
        pass

    @abstractmethod
    def _get_traffic_light(self, lane):
        """Get traffic light object from lane. Map-type specific."""
        pass

    @abstractmethod
    def _set_lane_on_intersection(self, lane, value: bool):
        """Set the on_intersection attribute for a lane. Map-type specific."""
        pass

    @abstractmethod
    def _set_lane_is_approaching(self, lane, value: bool):
        """Set the is_approaching attribute for a lane. Map-type specific."""
        pass

    @abstractmethod
    def _get_lane_on_intersection(self, lane) -> bool:
        """Get the on_intersection status of a lane. Map-type specific."""
        pass

    # Concrete methods using abstract methods
    def create_lane_dict(self):
        """Returns a dictionary mapping each lane's lane_id to the lane object."""
        self.lane_dict = {self._get_lane_id(lane): lane for lane in self.lanes.values()}
        return self.lane_dict

    def get_lane_successors_and_predecessors(self):
        """Returns dictionaries mapping each lane's lane_id to its successor and predecessor lane indices."""
        lane_successors = {}
        lane_predecessors = {}

        for lane in self.lanes.values():
            lane_id = self._get_lane_id(lane)
            lane_successors[lane_id] = self._get_lane_successors(lane)
            lane_predecessors[lane_id] = self._get_lane_predecessors(lane)

        self.lane_successors_dict = lane_successors
        self.lane_predecessors_dict = lane_predecessors
        return lane_successors, lane_predecessors

    def check_if_all_lanes_are_on_segment(self):
        """
        Checks if all lanes are on a segment.
        Returns:
            bool: True if all lanes are on a segment, False otherwise.
        """
        for lane in self.lanes.values():
            lane_id = self._get_lane_id(lane)
            if lane_id not in self.lane_segment_dict or self.lane_segment_dict[lane_id].segment is None:
                return False
        return True

check_if_all_lanes_are_on_segment()

Checks if all lanes are on a segment. Returns: bool: True if all lanes are on a segment, False otherwise.

Source code in omega_prime/mapsegment.py

def check_if_all_lanes_are_on_segment(self):
    """
    Checks if all lanes are on a segment.
    Returns:
        bool: True if all lanes are on a segment, False otherwise.
    """
    for lane in self.lanes.values():
        lane_id = self._get_lane_id(lane)
        if lane_id not in self.lane_segment_dict or self.lane_segment_dict[lane_id].segment is None:
            return False
    return True

create_lane_dict()

Returns a dictionary mapping each lane's lane_id to the lane object.

Source code in omega_prime/mapsegment.py

def create_lane_dict(self):
    """Returns a dictionary mapping each lane's lane_id to the lane object."""
    self.lane_dict = {self._get_lane_id(lane): lane for lane in self.lanes.values()}
    return self.lane_dict

get_lane_successors_and_predecessors()

Returns dictionaries mapping each lane's lane_id to its successor and predecessor lane indices.

Source code in omega_prime/mapsegment.py

def get_lane_successors_and_predecessors(self):
    """Returns dictionaries mapping each lane's lane_id to its successor and predecessor lane indices."""
    lane_successors = {}
    lane_predecessors = {}

    for lane in self.lanes.values():
        lane_id = self._get_lane_id(lane)
        lane_successors[lane_id] = self._get_lane_successors(lane)
        lane_predecessors[lane_id] = self._get_lane_predecessors(lane)

    self.lane_successors_dict = lane_successors
    self.lane_predecessors_dict = lane_predecessors
    return lane_successors, lane_predecessors

Segment

Bases: ABC

A class that represents a segment of the map

Source code in omega_prime/mapsegment.py

class Segment(ABC):
    """A class that represents a segment of the map"""

    def __init__(self, lanes, idx=None, concave_hull_ratio=0.3):
        self.lanes = lanes
        self.lane_ids = [self._get_lane_id(lane) for lane in lanes]
        self.trafficlights = []
        self.idx = idx
        self.concave_hull_ratio = concave_hull_ratio
        self.type = MapSegmentType.UNKNOWN

        # Cache polygon to avoid recomputing concave hull when lanes stay unchanged
        self._polygon_cache = None
        self._polygon_cache_key = None
        self._polygon_dirty = True
        self.polygon = self.create_segment_polygon()

    @abstractmethod
    def _get_lane_id(self, lane):
        """Extract lane ID from a lane object. Map-type specific."""
        pass

    @abstractmethod
    def _get_lane_geometry(self, lane) -> shapely.LineString:
        """Extract geometry from a lane object. Map-type specific."""
        pass

    @abstractmethod
    def set_trafficlight(self):
        """Set traffic lights for this segment. Map-type specific."""
        pass

    def _compute_polygon_key(self):
        return tuple((self._get_lane_id(lane), self._get_lane_geometry(lane).wkb) for lane in self.lanes)

    def _compute_segment_polygon(self):
        lane_centerline = [self._get_lane_geometry(lane) for lane in self.lanes]

        multi_centerline = shapely.geometry.MultiLineString(lane_centerline)
        combined = multi_centerline.buffer(0.1)
        combined = combined.simplify(0.1, preserve_topology=True)

        try:
            hull = shapely.concave_hull(combined, self.concave_hull_ratio)
            assert not hull.is_empty
        except (shapely.errors.GEOSException, AssertionError):
            hull = shapely.convex_hull(combined)
            assert not hull.is_empty
        return hull

    def _ensure_polygon(self, force=False):
        key = self._compute_polygon_key()
        if force or self._polygon_dirty or key != self._polygon_cache_key:
            self._polygon_cache = self._compute_segment_polygon()
            self._polygon_cache_key = key
            self._polygon_dirty = False
        return self._polygon_cache

    def get_center_point(self):
        "Returns the center point of the segment"
        return self.polygon.centroid.x, self.polygon.centroid.y

    def create_segment_polygon(self):
        "Create the Polygon of the Segment"
        return self._ensure_polygon()

    def update_polygon(self):
        "Updates the Polygon of the Segment"
        self._polygon_dirty = True
        self.polygon = self._ensure_polygon(force=True)

    def add_lane(self, lanes, update_polygon=True):
        """Adds a lane to the segment.
        If the lane is already in the segment, it will not be added again.

        Args:
            lane (list): A list of lane objects to be added to the segment.
        """
        for lane in lanes:
            if lane not in self.lanes:
                self.lanes.append(lane)
                self.lane_ids.append(self._get_lane_id(lane))

        if update_polygon:
            self.update_polygon()

        self.set_trafficlight()

    def get_timeinterval_on_segment(self, roaduser):
        """
        Gets a roadsegment as input as well as a roaduser trajectory.
        Returns the time interval of the roaduser on the segment.
        roaduser should be a np.array with (total_nanos, x, y)
        """
        if self.polygon:
            roaduser_points = [shapely.Point(x, y) for x, y in roaduser[:, 1:3]]
            roaduser_on_segment = np.array([self.polygon.contains(point) for point in roaduser_points])
            if roaduser_on_segment.any():
                indices = np.where(roaduser_on_segment)[0]
                return roaduser[indices[0], 0], roaduser[indices[-1], 0]
            else:
                return None
        else:
            return None

add_lane(lanes, update_polygon=True)

Adds a lane to the segment. If the lane is already in the segment, it will not be added again.

Parameters:

Name	Type	Description	Default
lane	list	A list of lane objects to be added to the segment.	required

Source code in omega_prime/mapsegment.py

def add_lane(self, lanes, update_polygon=True):
    """Adds a lane to the segment.
    If the lane is already in the segment, it will not be added again.

    Args:
        lane (list): A list of lane objects to be added to the segment.
    """
    for lane in lanes:
        if lane not in self.lanes:
            self.lanes.append(lane)
            self.lane_ids.append(self._get_lane_id(lane))

    if update_polygon:
        self.update_polygon()

    self.set_trafficlight()

create_segment_polygon()

Create the Polygon of the Segment

Source code in omega_prime/mapsegment.py

def create_segment_polygon(self):
    "Create the Polygon of the Segment"
    return self._ensure_polygon()

get_center_point()

Returns the center point of the segment

Source code in omega_prime/mapsegment.py

def get_center_point(self):
    "Returns the center point of the segment"
    return self.polygon.centroid.x, self.polygon.centroid.y

get_timeinterval_on_segment(roaduser)

Gets a roadsegment as input as well as a roaduser trajectory. Returns the time interval of the roaduser on the segment. roaduser should be a np.array with (total_nanos, x, y)

Source code in omega_prime/mapsegment.py

def get_timeinterval_on_segment(self, roaduser):
    """
    Gets a roadsegment as input as well as a roaduser trajectory.
    Returns the time interval of the roaduser on the segment.
    roaduser should be a np.array with (total_nanos, x, y)
    """
    if self.polygon:
        roaduser_points = [shapely.Point(x, y) for x, y in roaduser[:, 1:3]]
        roaduser_on_segment = np.array([self.polygon.contains(point) for point in roaduser_points])
        if roaduser_on_segment.any():
            indices = np.where(roaduser_on_segment)[0]
            return roaduser[indices[0], 0], roaduser[indices[-1], 0]
        else:
            return None
    else:
        return None

set_trafficlight() abstractmethod

Set traffic lights for this segment. Map-type specific.

Source code in omega_prime/mapsegment.py

@abstractmethod
def set_trafficlight(self):
    """Set traffic lights for this segment. Map-type specific."""
    pass

update_polygon()

Updates the Polygon of the Segment

Source code in omega_prime/mapsegment.py

def update_polygon(self):
    "Updates the Polygon of the Segment"
    self._polygon_dirty = True
    self.polygon = self._ensure_polygon(force=True)

ConnectionSegment

Bases: SegmentOsiCenterline

Source code in omega_prime/maposicenterlinesegmentation.py

class ConnectionSegment(SegmentOsiCenterline):
    def __init__(self, lanes, idx=None, concave_hull_ratio=0.3):
        super().__init__(lanes, idx, concave_hull_ratio=concave_hull_ratio)
        self.type = MapSegmentType.STRAIGHT
        self.intersection_idxs = set()

    def plot(self, output_plot: Path):
        """Plots the Connection segment

        Args:
            output_plot (Path): Path to the output directory.
        Returns:
            None
        """
        fig, ax = plt.subplots(1, 1)
        ax.set_aspect(1)
        # Add the index of the center line to the plot
        ax.set_title(f"Connection segment {self.idx}")
        for lane in self.lanes:
            ax.plot(*np.asarray(lane.centerline.xy)[:2], color="blue")
        for lane in self.lanes:
            m = int(np.ceil(len(lane.centerline.xy[0]) / 2))
            ax.annotate(
                lane.idx.lane_id,
                xy=(lane.centerline.xy[0][m], lane.centerline.xy[1][m]),
                fontsize=2,
                color="black",
                zorder=3,
            )
        # Plot the polygon into the intersection
        try:
            ax.plot(*self.polygon.exterior.xy, color="red", alpha=0.5, zorder=10)
        except:
            logging.warning(f"Connection {self.idx} has no polygon")
            pass
        ax.set_aspect(1)
        plt.title(f"Connection with {len(self.lanes)} lanes")
        plt.xlabel("X Coordinate")
        plt.ylabel("Y Coordinate")
        if output_plot is None:
            plt.show()
        elif isinstance(output_plot, Path) and output_plot.is_dir():
            output_plot.mkdir(parents=True, exist_ok=True)
            plt.savefig(output_plot / f"Connection{self.idx}.pdf")
        else:
            raise ValueError("output_plot must be a Path to a directory or None")
        plt.close()

plot(output_plot)

Plots the Connection segment

Parameters:

Name	Type	Description	Default
output_plot	Path	Path to the output directory.	required

Returns: None

Source code in omega_prime/maposicenterlinesegmentation.py

def plot(self, output_plot: Path):
    """Plots the Connection segment

    Args:
        output_plot (Path): Path to the output directory.
    Returns:
        None
    """
    fig, ax = plt.subplots(1, 1)
    ax.set_aspect(1)
    # Add the index of the center line to the plot
    ax.set_title(f"Connection segment {self.idx}")
    for lane in self.lanes:
        ax.plot(*np.asarray(lane.centerline.xy)[:2], color="blue")
    for lane in self.lanes:
        m = int(np.ceil(len(lane.centerline.xy[0]) / 2))
        ax.annotate(
            lane.idx.lane_id,
            xy=(lane.centerline.xy[0][m], lane.centerline.xy[1][m]),
            fontsize=2,
            color="black",
            zorder=3,
        )
    # Plot the polygon into the intersection
    try:
        ax.plot(*self.polygon.exterior.xy, color="red", alpha=0.5, zorder=10)
    except:
        logging.warning(f"Connection {self.idx} has no polygon")
        pass
    ax.set_aspect(1)
    plt.title(f"Connection with {len(self.lanes)} lanes")
    plt.xlabel("X Coordinate")
    plt.ylabel("Y Coordinate")
    if output_plot is None:
        plt.show()
    elif isinstance(output_plot, Path) and output_plot.is_dir():
        output_plot.mkdir(parents=True, exist_ok=True)
        plt.savefig(output_plot / f"Connection{self.idx}.pdf")
    else:
        raise ValueError("output_plot must be a Path to a directory or None")
    plt.close()

MapOsiCenterlineSegmentation

Bases: MapSegmentation

A class that identifies different segments on a OsiCenterline Map. Concrete implementation of MapSegmentation for OSI centerline maps.

Source code in omega_prime/maposicenterlinesegmentation.py

class MapOsiCenterlineSegmentation(MapSegmentation):
    """
    A class that identifies different segments on a OsiCenterline Map.
    Concrete implementation of MapSegmentation for OSI centerline maps.
    """

    def __init__(self, recording, lane_buffer=None, intersection_overlap_buffer=None, concave_hull_ratio=0.3):
        super().__init__(recording, concave_hull_ratio=concave_hull_ratio)
        self.locator = Locator.from_map(recording.map)
        self.isolated_connections = []
        self.G = None
        self.lane_buffer = lane_buffer if lane_buffer is not None else 0.3
        self.intersection_overlap_buffer = intersection_overlap_buffer if intersection_overlap_buffer is not None else 1
        self.do_combine_intersections = True

        for tl_state in recording.traffic_light_states.values():
            for tl in tl_state:
                if tl.id.value not in self.trafficlight_ids:
                    self.trafficlight[tl.id.value] = tl
                    self.trafficlight_ids.add(tl.id.value)

    # Implement abstract methods for OSI centerline maps
    def _get_lane_id(self, lane):
        """Extract lane ID from OSI centerline lane."""
        return lane.idx.lane_id

    def _get_lane_centerline(self, lane) -> shapely.LineString:
        """Extract centerline from OSI centerline lane."""
        return lane.centerline

    def _get_lane_successors(self, lane) -> list:
        """Extract successor IDs from OSI centerline lane."""
        return [succ_id.lane_id if hasattr(succ_id, "lane_id") else succ_id for succ_id in lane.successor_ids]

    def _get_lane_predecessors(self, lane) -> list:
        """Extract predecessor IDs from OSI centerline lane."""
        return [pred_id.lane_id if hasattr(pred_id, "lane_id") else pred_id for pred_id in lane.predecessor_ids]

    def _has_traffic_light(self, lane) -> bool:
        """Check if OSI centerline lane has traffic light."""
        return hasattr(lane, "trafficlight") and lane.trafficlight is not None

    def _get_traffic_light(self, lane):
        """Get traffic light object from OSI centerline lane."""
        return lane.trafficlight if self._has_traffic_light(lane) else None

    def _set_lane_on_intersection(self, lane, value: bool):
        """Set the on_intersection attribute for OSI centerline lane."""
        lane.on_intersection = value

    def _set_lane_is_approaching(self, lane, value: bool):
        """Set the is_approaching attribute for OSI centerline lane."""
        lane.is_approaching = value

    def _get_lane_on_intersection(self, lane) -> bool:
        """Get the on_intersection status of OSI centerline lane."""
        return lane.on_intersection if hasattr(lane, "on_intersection") else False

    def init_intersections(self):
        """
        Initializes the intersections in the map.
        Args:
            None
        Returns:
            None
        """
        self.create_lane_dict()
        self.get_lane_successors_and_predecessors()
        self.parallel_lane_dict = self.create_parallel_lane_dict()
        self.get_intersecting_lanes()
        self.set_lane_trafficlights()
        self.graph_intersection_detection()
        self.G = add_lanexy_to_graph(self.G, self.lanes)
        self.set_intersection_idx()

        if self.do_combine_intersections:
            self.create_lane_segment_dict()
            self.add_non_intersecting_lanes_to_intersection()
            self.combine_intersections()
            self.set_intersection_idx()

        self.create_intersection_dict()
        self.create_lane_segment_dict()
        self.find_isolated_connections()
        self.create_lane_segment_dict()
        self.check_if_all_lanes_are_on_segment()
        self.update_segment_ids()
        self.create_lane_segment_dict()
        self.update_road_ids()
        self.set_lane_intersection_relation()

        # from pathlib import Path
        # #Plot the graph G with x and y coordinates of the lanes
        # plot_graph(self.G , Path("/scenario-center-playground/scenarios/") / "graph_plot.pdf")

    def update_road_ids(self):
        """
        Updates the road_ids of the lane to the segment ID
        """
        updates_needed = []
        old_to_new_mapping = {}

        # First pass: identify what needs to be updated
        for lane_idx, lane in self.lanes.items():
            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:
                new_road_id = self.lane_segment_dict[lane_id].segment.idx
                if lane.idx.road_id != new_road_id:
                    new_idx = lane.idx._replace(road_id=new_road_id)
                    updates_needed.append((lane_idx, lane, new_idx))
                    old_to_new_mapping[lane_idx] = new_idx

        # Second pass: apply updates efficiently
        for old_idx, lane, new_idx in updates_needed:
            # Update the lane object in place
            lane.idx = new_idx

            # Only modify dictionary if the key actually changed
            if old_idx != new_idx:
                self.lanes[new_idx] = lane
                del self.lanes[old_idx]

        # Third pass: update all predecessor and successor references
        for lane in self.lanes.values():
            # Update predecessor references
            updated_predecessors = []
            for pred_id in lane.predecessor_ids:
                if pred_id in old_to_new_mapping:
                    updated_predecessors.append(old_to_new_mapping[pred_id])
                else:
                    updated_predecessors.append(pred_id)
            lane.predecessor_ids = updated_predecessors

            # Update successor references
            updated_successors = []
            for succ_id in lane.successor_ids:
                if succ_id in old_to_new_mapping:
                    updated_successors.append(old_to_new_mapping[succ_id])
                else:
                    updated_successors.append(succ_id)
            lane.successor_ids = updated_successors

        # Fourth pass: update internal dictionaries that track relationships
        self.lane_dict = {lane.idx.lane_id: lane for lane in self.lanes.values()}
        self.get_lane_successors_and_predecessors()

    def update_segment_ids(self):
        "Updates the segment IDs of the map segmentation"
        self.segments = self.intersections + self.isolated_connections
        for i, segment in enumerate(self.segments):
            segment.idx = i
            segment.set_trafficlight()

    def create_parallel_lane_dict(self):
        """
        Creates a dictionary mapping each lane's lane_id to the lane ids which are parallel to it
        Args:
            None
        Returns:
            dict: A dictionary mapping each lane's lane_id to the lane ids which are parallel to it.
        """
        lane_dict = {lane.idx.lane_id: [] for lane in self.lanes.values()}

        # Precompute lane directions for faster comparisons
        lane_directions = {}
        lane_centerlines = []
        lane_ids = []

        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)
            lane_centerlines.append(lane.centerline)
            lane_ids.append(lane.idx.lane_id)

        if not lane_centerlines:
            return lane_dict

        # Use original centerlines for spatial index, buffer only when needed
        tree = STRtree(lane_centerlines)

        for i, lane in enumerate(self.lanes.values()):
            lane_id = lane.idx.lane_id

            # Create buffer only when querying, not storing it
            buffer_geom = lane.centerline.buffer(10)
            candidates = tree.query(buffer_geom)

            # Clear the buffer immediately after use
            del buffer_geom

            for idx in candidates:
                other_lane_id = lane_ids[idx]
                if other_lane_id == lane_id:
                    continue

                # Compare directions using dot product
                dir1 = lane_directions[lane_id]
                dir2 = lane_directions[other_lane_id]
                dot_product = np.clip(np.abs(np.dot(dir1, dir2)), -1.0, 1.0)
                angle_deg = np.degrees(np.arccos(dot_product))

                if angle_deg < 10:
                    lane_dict[lane_id].append(other_lane_id)

        return lane_dict

    def trajectory_segment_detection(self, trajectory):
        """
        Splits a trajectory into segments based on the lane it is located on

        Args:
            trajectory (np.ndarray): A NumPy array of shape (n, 3) representing the trajectory, where each row is a (frame, x, y) coordinate.

        Returns:
            list: A list of tuples, where each tuple contains a segment of the trajectory and the segment it intersects with.
        """
        segments = []
        current_segment = []
        xy = trajectory[:, 1:3]  # Extract x and y coordinates
        sts = self.locator.xys2sts(xy)
        lane_ids = sts["roadlane_id"].to_numpy()
        segment_idx = [self.lane_segment_dict[lane_id.lane_id].segment.idx for lane_id in lane_ids]

        trajectory = np.column_stack((trajectory[:, 0], trajectory[:, 1], trajectory[:, 2], lane_ids, segment_idx))

        # Create spatial index for intersection polygons
        intersection_polygons = []
        intersection_ids = []
        buffer = 5

        for segment in self.segments:
            if segment.type == MapSegmentType.JUNCTION and hasattr(segment, "polygon"):
                intersection_polygons.append(segment.polygon.buffer(buffer))
                intersection_ids.append(segment.idx)

        if intersection_polygons:
            # Use spatial index for efficient intersection queries
            tree = STRtree(intersection_polygons)

            # Process points in batches for better performance
            for i, (frame, x, y, _, _) in enumerate(trajectory):
                point = Point(x, y)

                # Query spatial index instead of checking all polygons
                candidates = tree.query(point)

                for idx in candidates:
                    if intersection_polygons[idx].contains(point):
                        trajectory[i, 4] = intersection_ids[idx]
                        break

        # Rest of the method for creating segments
        prev_seg_id = -1
        for i, (frame, x, y, _, segment_idx) in enumerate(trajectory):
            if prev_seg_id == segment_idx:
                current_segment.append((frame, x, y))
            else:
                if current_segment:
                    segments.append((np.array(current_segment), self.segments[prev_seg_id]))
                current_segment = [(frame, x, y)]
                prev_seg_id = segment_idx

        if current_segment:
            segments.append((np.array(current_segment), self.segments[prev_seg_id]))

        return segments

    def get_intersecting_lanes(self, buffer: float = None):
        """
        Returns a dictionary mapping each lane's lane_id to an array of lane ids it intersects with.

        Args:
            lanes (list): Array of lane objects, each with an `idx` and `centerline` attribute.

        Returns:
            dict: A dictionary where keys are lane ids and values are arrays of intersecting lane ids.
        """
        if buffer is None:
            buffer = self.lane_buffer

        # Create spatial index directly from centerlines
        lane_centerlines = []
        lane_ids = []

        for lane in self.lanes.values():
            lane_centerlines.append(lane.centerline)
            lane_ids.append(lane.idx.lane_id)

        if not lane_centerlines:
            self.intersecting_lanes_dict = {}
            return {}

        tree = STRtree(lane_centerlines)

        # Pre-compute lane relationships for faster lookup
        successors_set = {lane_id: set(successors) for lane_id, successors in self.lane_successors_dict.items()}
        predecessors_set = {lane_id: set(predecessors) for lane_id, predecessors in self.lane_predecessors_dict.items()}
        parallel_set = {lane_id: set(parallel) for lane_id, parallel in self.parallel_lane_dict.items()}

        intersecting_lanes = {}
        for i, lane in enumerate(self.lanes.values()):
            lane_id = lane.idx.lane_id

            # Query spatial index with buffered geometry
            buffered_centerline = lane.centerline.buffer(buffer)
            candidates = tree.query(buffered_centerline)

            intersecting_lanes[lane_id] = []
            for idx in candidates:
                candidate_id = lane_ids[idx]
                if (
                    candidate_id != lane_id
                    and candidate_id not in successors_set[lane_id]
                    and candidate_id not in predecessors_set[lane_id]
                    and candidate_id not in parallel_set[lane_id]
                    and buffered_centerline.intersects(lane_centerlines[idx])
                ):
                    # Only buffer and test intersection for valid candidates
                    intersecting_lanes[lane_id].append(candidate_id)

        self.intersecting_lanes_dict = intersecting_lanes
        return intersecting_lanes

    def graph_intersection_detection(self):
        """
        Detects intersections in a graph of lanes based on their intersections, successors, and predecessors.

        Args:
            lane_dict (dict): A dictionary where keys are lane indices and values are arrays of intersecting lane indices.
            lane_successors (dict): A dictionary where keys are lane indices and values are arrays of successor lane indices.
            lane_predecessors (dict): A dictionary where keys are lane indices and values are arrays of predecessor lane indices.

        Returns:
            list: A list of intersections, where each intersection is a set of lane indices.
        """
        # Create a Graph using networkx
        G = nx.Graph()

        # Add nodes and edges to the graph. If a lane has a intersection, add the lanes as nodes and the intersection as an edge
        # Add edges directly (nodes are added automatically)
        for lane_id, intersecting_lanes in self.intersecting_lanes_dict.items():
            G.add_edges_from((lane_id, other_lane) for other_lane in intersecting_lanes)

        intersections = []
        for inter in nx.connected_components(G):
            # Convert lane_ids back to lane objects
            intersection_lanes = [self.lane_dict[i] for i in inter]
            intersection = Intersection(intersection_lanes, concave_hull_ratio=self.concave_hull_ratio)
            intersections.append(intersection)

        self.intersections = intersections
        self.G = G
        return intersections, G

    def combine_intersections(self):
        """A function that revieves a list with idx [[1,2] , [4,5,6] , ...] of intersections that need to be combined.
        It will combine all those intersections and will then update all intersections in the map_segmentation class.

        Args:
            intersection_list (list): A list of lists, where each inner list contains the indices of intersections to be combined.

        Returns:
            None
        """

        # Check for intersections that can be combined:
        combined_intersections = []

        # Create spatial index of all intersection polygons
        for intersection in self.intersections:
            if (
                not hasattr(intersection, "_buffered_polygon")
                or intersection._buffer_value != self.intersection_overlap_buffer
            ):
                intersection._buffered_polygon = intersection.polygon.buffer(self.intersection_overlap_buffer)
                intersection._buffer_value = self.intersection_overlap_buffer

        polygons = [intersection._buffered_polygon for intersection in self.intersections]
        if polygons:
            tree = STRtree(polygons)

            # Find overlapping intersections efficiently
            for i, intersection in enumerate(self.intersections):
                buffered_poly = polygons[i]
                candidates = tree.query(buffered_poly)

                for j in candidates:
                    if i != j and buffered_poly.intersects(polygons[j]):
                        combined_intersections.append([i, j])
        final_combined = self.find_resulting_intersections(combined_intersections)
        new_intersections = []
        visited = set()

        for combination in final_combined:
            combined_lanes = []
            for idx in combination:
                if idx not in visited:
                    visited.add(idx)
                    combined_lanes.extend(self.intersections[idx].lanes)

            new_intersections.append(Intersection(combined_lanes, concave_hull_ratio=self.concave_hull_ratio))

        # Add unvisited intersections
        for i, intersection in enumerate(self.intersections):
            if i not in visited:
                new_intersections.append(intersection)

        self.intersections = new_intersections

    def intersections_overlap(self, intersection1, intersection2, buffer: float = None):
        """
        Check if two intersections overlap.

        Args:
            intersection1 (Intersection): The first intersection object.
            intersection2 (Intersection): The second intersection object.

        Returns:
            bool: True if the intersections overlap, False otherwise.
        """
        if buffer is None:
            buffer = self.intersection_overlap_buffer

        # Use cached buffers if available
        if not hasattr(intersection1, "_buffered_polygon") or intersection1._buffer_value != buffer:
            intersection1._buffered_polygon = intersection1.polygon.buffer(buffer)
            intersection1._buffer_value = buffer

        if not hasattr(intersection2, "_buffered_polygon") or intersection2._buffer_value != buffer:
            intersection2._buffered_polygon = intersection2.polygon.buffer(buffer)
            intersection2._buffer_value = buffer

        return intersection1.polygon.buffer(buffer).intersects(intersection2.polygon.buffer(buffer))

    def combine_intersection_on_polygon(self, intersection1, intersection2):
        """
        Combine two intersections into one if they overlap.
        Args:
            intersection1 (Intersection): The first intersection object.
            intersection2 (Intersection): The second intersection object.
        Returns:
            Intersection: The combined intersection object if they overlap, None otherwise.
        """
        if self.intersections_overlap(intersection1, intersection2):
            # Create a new intersection object with the lanes from both intersections
            combined_intersection = Intersection(
                intersection1.lanes + intersection2.lanes, concave_hull_ratio=self.concave_hull_ratio
            )

            return combined_intersection
        else:
            return None

    def find_resulting_intersections(self, intersection_pairs):
        G = nx.Graph()
        G.add_edges_from(intersection_pairs)
        return [list(component) for component in nx.connected_components(G)]

    def set_intersection_idx(self):
        """
        Sets the index for each intersection in the list of intersections.
        Args:
            None
        Returns:
            None
        """
        for i, intersection in enumerate(self.intersections + self.isolated_connections):
            intersection.idx = i

    def create_intersection_dict(self):
        """Creats a dictionary where the key is the intersection id and the value is the intersection object.
        Args:
            None
        Returns:
            None
        """
        intersection_dict = {}
        for i, intersection in enumerate(self.intersections):
            intersection_dict[intersection.idx] = intersection
        self.intersection_dict = intersection_dict
        return intersection_dict

    def add_non_intersecting_lanes_to_intersection(self):
        """Add all lanes that are within the intersection polygon to the intersection object.
        Args:
            None
        Returns:
            None
        """
        for intersection in self.intersections:
            intersection.update_polygon()

            # Collect all lanes to add before modifying the intersection
            lanes_to_add = []
            buffered_polygon = intersection.polygon.buffer(self.lane_buffer)

            for lane in self.lanes.values():
                lane_id = lane.idx.lane_id
                if (
                    lane_id not in intersection.lane_ids
                    and self.lane_segment_dict[lane_id].segment is None
                    and buffered_polygon.contains(lane.centerline)
                ):
                    lanes_to_add.append(lane)

            # Add all lanes at once and update polygon only once
            if lanes_to_add:
                intersection.add_lane(lanes=lanes_to_add, update_polygon=True)  # Assuming bulk add method

    def create_lane_segment_dict(self):
        """
        Create a dictionary mapping lane IDs to their segment information.
        Args:
            None
        Returns:
            lane_segment_dict (dict): A dictionary mapping lane IDs to their segment information.
        """
        segment_name = nt("SegmentName", ["lane_id", "segment_idx", "segment"])
        segment_list = self.intersections + self.isolated_connections

        # Initialize with None values more efficiently
        lane_segment_dict = {lane_id: segment_name(lane_id, None, None) for lane_id in self.lane_dict.keys()}

        for segment in segment_list:
            for lane in segment.lanes:
                lane_id = lane.idx.lane_id

                # Single lookup with caching
                current_entry = lane_segment_dict.get(lane_id)
                if current_entry is None:
                    continue

                if current_entry.segment is None:
                    # Lane not assigned to any segment yet
                    lane_segment_dict[lane_id] = segment_name(lane_id, segment.idx, segment)
                elif current_entry.segment_idx != segment.idx:
                    # Conflict: lane already assigned to different segment
                    logger.warning(
                        f"Lane {lane_id} already in segment {current_entry.segment_idx}, "
                        f"cannot assign to segment {segment.idx}"
                    )

        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.
        Args:
            None
        Returns:
            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.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)
                for successor in self.lane_successors_dict[lane_id]:
                    if successor not in self.lane_segment_dict or self.lane_segment_dict[successor].segment is None:
                        G.add_edge(lane_id, successor)
                for predecessor in self.lane_predecessors_dict[lane_id]:
                    if predecessor not in self.lane_segment_dict or self.lane_segment_dict[predecessor].segment is None:
                        G.add_edge(lane_id, predecessor)
        return G

    def find_isolated_connections(self):
        """Find all isolated strings of connections in the graph. Then Check if any of those lanes would be part of an intersection.
        Args:
            None
        Returns:
            isolated_connections (list): A list of lists, where each inner list contains the indices of lanes that are part of an isolated connection.
        """
        G = self.create_non_intersecting_lane_graph()
        new_connections = []
        segment_name_type = type(next(iter(self.lane_segment_dict.values())))

        # Classify each component before constructing ConnectionSegment objects.
        for component in nx.connected_components(G):
            if not component:
                continue

            component_ids = list(component)
            pre = False
            suc = False
            intersection_idxs = set()

            for lane_id in component_ids:
                for successor in self.lane_successors_dict[lane_id]:
                    if successor in self.lane_segment_dict and self.lane_segment_dict[successor].segment is not None:
                        intersection_idxs.add(self.lane_segment_dict[successor].segment_idx)
                        suc = True
                for predecessor in self.lane_predecessors_dict[lane_id]:
                    if (
                        predecessor in self.lane_segment_dict
                        and self.lane_segment_dict[predecessor].segment is not None
                    ):
                        intersection_idxs.add(self.lane_segment_dict[predecessor].segment_idx)
                        pre = True

            if len(intersection_idxs) == 1 and pre and suc:
                # Single bordering intersection on both ends: absorb the component into it.
                inter = self.intersection_dict[list(intersection_idxs)[0]]
                for lane_id in component_ids:
                    inter.lanes.append(self.lane_dict[lane_id])
                    inter.lane_ids.append(lane_id)
                    # Keep lane_segment_dict current so subsequent components see these lanes as already assigned.
                    self.lane_segment_dict[lane_id] = segment_name_type(lane_id, inter.idx, inter)
                inter.update_polygon()
            else:
                connection = ConnectionSegment(
                    [self.lane_dict[i] for i in component_ids], concave_hull_ratio=self.concave_hull_ratio
                )
                connection.intersection_idxs = intersection_idxs
                new_connections.append(connection)

        isolated_connections = new_connections
        # Create ConnectionSegment for all lanes, that are on multiple intersections:

        isolated_connections = self.combine_isolated_connections(isolated_connections)
        return isolated_connections

    def combine_isolated_connections(self, isolated_connections):
        """Check if any of the isolated connections are connecting the same intersections.
        If yes, then combine them into one connection.
        Args:
            isolated_connections (list): A list of ConnectionSegment objects representing isolated connections.
        Returns:
            isolated_connections (list): A list of ConnectionSegment objects representing the combined isolated connections.
        """
        if not isolated_connections:
            return []

        # Group connections by their intersection indices for efficient comparison
        connections_by_intersections = {}
        for i, connection in enumerate(isolated_connections):
            key = frozenset(connection.intersection_idxs)
            if key not in connections_by_intersections:
                connections_by_intersections[key] = []
            connections_by_intersections[key].append(i)

        combined_connections = []

        # Process each group of connections with same intersection indices
        for intersection_set, connection_indices in connections_by_intersections.items():
            if len(connection_indices) > 1:
                if len(intersection_set) > 1:
                    # Multiple intersections: combine all connections
                    for i in range(len(connection_indices)):
                        for j in range(i + 1, len(connection_indices)):
                            combined_connections.append([connection_indices[i], connection_indices[j]])
                else:
                    # Single intersection: check distance
                    connections_to_check = [isolated_connections[idx] for idx in connection_indices]
                    for i in range(len(connections_to_check)):
                        for j in range(i + 1, len(connections_to_check)):
                            if connections_to_check[i].polygon.distance(connections_to_check[j].polygon) < 5:
                                combined_connections.append([connection_indices[i], connection_indices[j]])

        # Rest of the method remains the same
        final_combined = self.find_resulting_intersections(combined_connections)
        new_connections = []
        visited = set()

        for combination in final_combined:
            combined_lanes = []
            for idx in combination:
                if idx not in visited:
                    visited.add(idx)
                    combined_lanes.extend(isolated_connections[idx].lanes)
            new_connections.append(ConnectionSegment(combined_lanes, concave_hull_ratio=self.concave_hull_ratio))

        # Add unvisited connections
        for i, connection in enumerate(isolated_connections):
            if i not in visited:
                new_connections.append(connection)

        self.isolated_connections = new_connections
        return self.isolated_connections

    def set_lane_intersection_relation(self):
        """
        Sets the attribute lane.is_approaching true if the lane is connecting to an intersection.
        Sets the attribute lane.on_intersection true if the lane is part of an intersection.
        """
        for lane in self.lanes.values():
            self._set_lane_on_intersection(lane, False)
            self._set_lane_is_approaching(lane, False)

        # Process intersection lanes and their predecessors efficiently
        for intersection in self.intersections:
            # Mark intersection lanes
            for lane in intersection.lanes:
                lane_id = self._get_lane_id(lane)
                if lane_id in self.lanes:
                    self._set_lane_on_intersection(self.lanes[lane_id], True)

                # Process predecessors for each lane in the intersection
                for predecessor_id in self._get_lane_predecessors(lane):
                    if predecessor_id in self.lane_dict:
                        predecessor = self.lane_dict[predecessor_id]
                        if not self._get_lane_on_intersection(predecessor):
                            self._set_lane_is_approaching(predecessor, True)

    def set_lane_trafficlights(self):
        """
        Sets the traffic lights for each lane of the map.
        """
        # Create spatial index for lane centerlines
        lane_centerlines = [lane.centerline for lane in self.lanes.values()]
        lane_objects = list(self.lanes.values())

        tree = STRtree(lane_centerlines)

        for tl_idx in self.trafficlight:
            traffic_light_found = False

            # Create traffic light position point
            tl_point = Point(self.trafficlight[tl_idx].base.position.x, self.trafficlight[tl_idx].base.position.y)

            # Use spatial index to find candidate lanes
            candidates = tree.nearest(tl_point)

            if candidates:
                lane = lane_objects[candidates]
                lane.traffic_light = self.trafficlight[tl_idx]
                traffic_light_found = True

            if not traffic_light_found:
                logger.warning(f"Traffic light {self.trafficlight[tl_idx].id} not found in any lane")

    def plot(
        self,
        output_plot: Path = None,
        trajectory=None,
        plot_lane_ids=False,
        plot_intersection_polygons=False,
        plot_connection_polygons=False,
    ):
        """
        Plots the intersections and saves the plot to the specified output path.
        A Trajectory can be given to plot it on the map. The Trajectory should be a numpy array of shape (n,3) where each row is (frame, x, y)
        Args:
            output_plot (Path): Path to a folder where the plot will be saved. If None, the plot will be shown instead.
            trajectory (numpy.ndarray): The trajectory to be plotted. If None, no trajectory will be plotted.
            plot_lane_ids (bool): Whether to plot lane IDs on the map.
            plot_intersection_polygons (bool): Whether to plot intersection polygons.
            plot_connection_polygons (bool): Whether to plot connection polygons.
        Returns:
            None
        """
        # Plot the map by plotting all the centerlines:
        fig, ax = plt.subplots(1, 1)
        ax.set_aspect(1)

        for lane in self.lanes.values():
            c = "blue"
            if lane.on_intersection:
                c = "green"
            elif lane.is_approaching:
                c = "orange"
            else:
                c = "black"
            ax.plot(*lane.centerline.xy, color=c, alpha=0.3, zorder=-10)

        if plot_lane_ids:
            lane_midpoints = [
                (lane.idx, lane.centerline.interpolate(0.5, normalized=True)) for lane in self.lanes.values()
            ]
            for lane_id, midpoint in lane_midpoints:
                ax.annotate(lane_id, xy=(midpoint.x, midpoint.y), fontsize=2, color="black")

        for inter in self.intersections:
            ax.annotate(inter.idx, xy=inter.get_center_point(), fontsize=2, color="black")

            if plot_intersection_polygons:
                # Plot the polygon into the intersection
                inter.update_polygon()
                ax.plot(*inter.polygon.exterior.xy, color="red", alpha=0.5, zorder=10)

        for combi in self.isolated_connections:
            ax.annotate(combi.idx, xy=combi.get_center_point(), fontsize=2, color="black")
            # Plot the polygon into the intersection
            if plot_connection_polygons:
                combi.update_polygon()
                try:
                    ax.plot(*combi.polygon.exterior.xy, color="blue", alpha=0.5, zorder=10)
                except:
                    logger.warning(f"Connection {combi.idx} has no polygon")
                    pass

        for tl_idx in self.trafficlight:
            position = shapely.Point(
                self.trafficlight[tl_idx].base.position.x, self.trafficlight[tl_idx].base.position.y
            )
            ax.plot(
                position.x,
                position.y,
                marker="o",
                color="red",
                markersize=2,
                label=f"Traffic Light {self.trafficlight[tl_idx].id}",
            )

        # Plot the trajectory if it is given
        if trajectory is not None:
            plt.plot(
                trajectory[:, 1],
                trajectory[:, 2],
                color="yellow",
                alpha=0.8,
                linewidth=3,
                label="Host Vehicle Trajectory",
            )

            # Mark start and end points
            plt.plot(trajectory[0, 1], trajectory[0, 2], "go", markersize=10, label="Start")
            plt.plot(trajectory[-1, 1], trajectory[-1, 2], "ro", markersize=10, label="End")

        ax.set_xlim(*ax.get_xlim())
        ax.set_ylim(*ax.get_ylim())
        plt.title("Map with Intersections")
        plt.xlabel("X Coordinate (m)", fontsize=12)
        plt.ylabel("Y Coordinate (m)", fontsize=12)
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.axis("equal")
        if output_plot is None:
            plt.show()
        else:
            if isinstance(output_plot, Path):
                output_plot.mkdir(parents=True, exist_ok=True)
                plt.savefig(output_plot / "Map_with_Intersection.pdf")
            else:
                isinstance(output_plot, str)
                output_path = Path(output_plot)
                if output_path.is_dir():
                    output_path.mkdir(parents=True, exist_ok=True)
                    plt.savefig(output_path / "Map_with_Intersection.pdf")
                elif output_path.suffix == ".pdf":
                    output_path.parent.mkdir(parents=True, exist_ok=True)
                    plt.savefig(output_path)
        plt.close()

    def plot_intersections(self, output_plot: Path):
        """
        Plots all intersections and saves them to the output path.
        Args:
            output_plot (Path): Path to a folder where the plots will be saved.
        Returns:
            None
        """
        for i, intersection in enumerate(self.intersections):
            intersection.plot(output_plot)
        for i, connection in enumerate(self.isolated_connections):
            connection.plot(output_plot)

add_non_intersecting_lanes_to_intersection()

Add all lanes that are within the intersection polygon to the intersection object. Args: None Returns: None

Source code in omega_prime/maposicenterlinesegmentation.py

def add_non_intersecting_lanes_to_intersection(self):
    """Add all lanes that are within the intersection polygon to the intersection object.
    Args:
        None
    Returns:
        None
    """
    for intersection in self.intersections:
        intersection.update_polygon()

        # Collect all lanes to add before modifying the intersection
        lanes_to_add = []
        buffered_polygon = intersection.polygon.buffer(self.lane_buffer)

        for lane in self.lanes.values():
            lane_id = lane.idx.lane_id
            if (
                lane_id not in intersection.lane_ids
                and self.lane_segment_dict[lane_id].segment is None
                and buffered_polygon.contains(lane.centerline)
            ):
                lanes_to_add.append(lane)

        # Add all lanes at once and update polygon only once
        if lanes_to_add:
            intersection.add_lane(lanes=lanes_to_add, update_polygon=True)  # Assuming bulk add method

combine_intersection_on_polygon(intersection1, intersection2)

Combine two intersections into one if they overlap. Args: intersection1 (Intersection): The first intersection object. intersection2 (Intersection): The second intersection object. Returns: Intersection: The combined intersection object if they overlap, None otherwise.

Source code in omega_prime/maposicenterlinesegmentation.py

def combine_intersection_on_polygon(self, intersection1, intersection2):
    """
    Combine two intersections into one if they overlap.
    Args:
        intersection1 (Intersection): The first intersection object.
        intersection2 (Intersection): The second intersection object.
    Returns:
        Intersection: The combined intersection object if they overlap, None otherwise.
    """
    if self.intersections_overlap(intersection1, intersection2):
        # Create a new intersection object with the lanes from both intersections
        combined_intersection = Intersection(
            intersection1.lanes + intersection2.lanes, concave_hull_ratio=self.concave_hull_ratio
        )

        return combined_intersection
    else:
        return None

combine_intersections()

A function that revieves a list with idx [[1,2] , [4,5,6] , ...] of intersections that need to be combined. It will combine all those intersections and will then update all intersections in the map_segmentation class.

Parameters:

Name	Type	Description	Default
intersection_list	list	A list of lists, where each inner list contains the indices of intersections to be combined.	required

Returns:

Type	Description
	None

Source code in omega_prime/maposicenterlinesegmentation.py

def combine_intersections(self):
    """A function that revieves a list with idx [[1,2] , [4,5,6] , ...] of intersections that need to be combined.
    It will combine all those intersections and will then update all intersections in the map_segmentation class.

    Args:
        intersection_list (list): A list of lists, where each inner list contains the indices of intersections to be combined.

    Returns:
        None
    """

    # Check for intersections that can be combined:
    combined_intersections = []

    # Create spatial index of all intersection polygons
    for intersection in self.intersections:
        if (
            not hasattr(intersection, "_buffered_polygon")
            or intersection._buffer_value != self.intersection_overlap_buffer
        ):
            intersection._buffered_polygon = intersection.polygon.buffer(self.intersection_overlap_buffer)
            intersection._buffer_value = self.intersection_overlap_buffer

    polygons = [intersection._buffered_polygon for intersection in self.intersections]
    if polygons:
        tree = STRtree(polygons)

        # Find overlapping intersections efficiently
        for i, intersection in enumerate(self.intersections):
            buffered_poly = polygons[i]
            candidates = tree.query(buffered_poly)

            for j in candidates:
                if i != j and buffered_poly.intersects(polygons[j]):
                    combined_intersections.append([i, j])
    final_combined = self.find_resulting_intersections(combined_intersections)
    new_intersections = []
    visited = set()

    for combination in final_combined:
        combined_lanes = []
        for idx in combination:
            if idx not in visited:
                visited.add(idx)
                combined_lanes.extend(self.intersections[idx].lanes)

        new_intersections.append(Intersection(combined_lanes, concave_hull_ratio=self.concave_hull_ratio))

    # Add unvisited intersections
    for i, intersection in enumerate(self.intersections):
        if i not in visited:
            new_intersections.append(intersection)

    self.intersections = new_intersections

combine_isolated_connections(isolated_connections)

Check if any of the isolated connections are connecting the same intersections. If yes, then combine them into one connection. Args: isolated_connections (list): A list of ConnectionSegment objects representing isolated connections. Returns: isolated_connections (list): A list of ConnectionSegment objects representing the combined isolated connections.

Source code in omega_prime/maposicenterlinesegmentation.py

def combine_isolated_connections(self, isolated_connections):
    """Check if any of the isolated connections are connecting the same intersections.
    If yes, then combine them into one connection.
    Args:
        isolated_connections (list): A list of ConnectionSegment objects representing isolated connections.
    Returns:
        isolated_connections (list): A list of ConnectionSegment objects representing the combined isolated connections.
    """
    if not isolated_connections:
        return []

    # Group connections by their intersection indices for efficient comparison
    connections_by_intersections = {}
    for i, connection in enumerate(isolated_connections):
        key = frozenset(connection.intersection_idxs)
        if key not in connections_by_intersections:
            connections_by_intersections[key] = []
        connections_by_intersections[key].append(i)

    combined_connections = []

    # Process each group of connections with same intersection indices
    for intersection_set, connection_indices in connections_by_intersections.items():
        if len(connection_indices) > 1:
            if len(intersection_set) > 1:
                # Multiple intersections: combine all connections
                for i in range(len(connection_indices)):
                    for j in range(i + 1, len(connection_indices)):
                        combined_connections.append([connection_indices[i], connection_indices[j]])
            else:
                # Single intersection: check distance
                connections_to_check = [isolated_connections[idx] for idx in connection_indices]
                for i in range(len(connections_to_check)):
                    for j in range(i + 1, len(connections_to_check)):
                        if connections_to_check[i].polygon.distance(connections_to_check[j].polygon) < 5:
                            combined_connections.append([connection_indices[i], connection_indices[j]])

    # Rest of the method remains the same
    final_combined = self.find_resulting_intersections(combined_connections)
    new_connections = []
    visited = set()

    for combination in final_combined:
        combined_lanes = []
        for idx in combination:
            if idx not in visited:
                visited.add(idx)
                combined_lanes.extend(isolated_connections[idx].lanes)
        new_connections.append(ConnectionSegment(combined_lanes, concave_hull_ratio=self.concave_hull_ratio))

    # Add unvisited connections
    for i, connection in enumerate(isolated_connections):
        if i not in visited:
            new_connections.append(connection)

    self.isolated_connections = new_connections
    return self.isolated_connections

create_intersection_dict()

Creats a dictionary where the key is the intersection id and the value is the intersection object. Args: None Returns: None

Source code in omega_prime/maposicenterlinesegmentation.py

def create_intersection_dict(self):
    """Creats a dictionary where the key is the intersection id and the value is the intersection object.
    Args:
        None
    Returns:
        None
    """
    intersection_dict = {}
    for i, intersection in enumerate(self.intersections):
        intersection_dict[intersection.idx] = intersection
    self.intersection_dict = intersection_dict
    return intersection_dict

create_lane_segment_dict()

Create a dictionary mapping lane IDs to their segment information. Args: None Returns: lane_segment_dict (dict): A dictionary mapping lane IDs to their segment information.

Source code in omega_prime/maposicenterlinesegmentation.py

def create_lane_segment_dict(self):
    """
    Create a dictionary mapping lane IDs to their segment information.
    Args:
        None
    Returns:
        lane_segment_dict (dict): A dictionary mapping lane IDs to their segment information.
    """
    segment_name = nt("SegmentName", ["lane_id", "segment_idx", "segment"])
    segment_list = self.intersections + self.isolated_connections

    # Initialize with None values more efficiently
    lane_segment_dict = {lane_id: segment_name(lane_id, None, None) for lane_id in self.lane_dict.keys()}

    for segment in segment_list:
        for lane in segment.lanes:
            lane_id = lane.idx.lane_id

            # Single lookup with caching
            current_entry = lane_segment_dict.get(lane_id)
            if current_entry is None:
                continue

            if current_entry.segment is None:
                # Lane not assigned to any segment yet
                lane_segment_dict[lane_id] = segment_name(lane_id, segment.idx, segment)
            elif current_entry.segment_idx != segment.idx:
                # Conflict: lane already assigned to different segment
                logger.warning(
                    f"Lane {lane_id} already in segment {current_entry.segment_idx}, "
                    f"cannot assign to segment {segment.idx}"
                )

    self.lane_segment_dict = lane_segment_dict

create_non_intersecting_lane_graph()

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. Args: None Returns: G (networkx.Graph): A graph with each lane as a node and the edges are the successors and predecessors of the lanes.

Source code in omega_prime/maposicenterlinesegmentation.py

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.
    Args:
        None
    Returns:
        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.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)
            for successor in self.lane_successors_dict[lane_id]:
                if successor not in self.lane_segment_dict or self.lane_segment_dict[successor].segment is None:
                    G.add_edge(lane_id, successor)
            for predecessor in self.lane_predecessors_dict[lane_id]:
                if predecessor not in self.lane_segment_dict or self.lane_segment_dict[predecessor].segment is None:
                    G.add_edge(lane_id, predecessor)
    return G

create_parallel_lane_dict()

Creates a dictionary mapping each lane's lane_id to the lane ids which are parallel to it Args: None Returns: dict: A dictionary mapping each lane's lane_id to the lane ids which are parallel to it.

Source code in omega_prime/maposicenterlinesegmentation.py

def create_parallel_lane_dict(self):
    """
    Creates a dictionary mapping each lane's lane_id to the lane ids which are parallel to it
    Args:
        None
    Returns:
        dict: A dictionary mapping each lane's lane_id to the lane ids which are parallel to it.
    """
    lane_dict = {lane.idx.lane_id: [] for lane in self.lanes.values()}

    # Precompute lane directions for faster comparisons
    lane_directions = {}
    lane_centerlines = []
    lane_ids = []

    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)
        lane_centerlines.append(lane.centerline)
        lane_ids.append(lane.idx.lane_id)

    if not lane_centerlines:
        return lane_dict

    # Use original centerlines for spatial index, buffer only when needed
    tree = STRtree(lane_centerlines)

    for i, lane in enumerate(self.lanes.values()):
        lane_id = lane.idx.lane_id

        # Create buffer only when querying, not storing it
        buffer_geom = lane.centerline.buffer(10)
        candidates = tree.query(buffer_geom)

        # Clear the buffer immediately after use
        del buffer_geom

        for idx in candidates:
            other_lane_id = lane_ids[idx]
            if other_lane_id == lane_id:
                continue

            # Compare directions using dot product
            dir1 = lane_directions[lane_id]
            dir2 = lane_directions[other_lane_id]
            dot_product = np.clip(np.abs(np.dot(dir1, dir2)), -1.0, 1.0)
            angle_deg = np.degrees(np.arccos(dot_product))

            if angle_deg < 10:
                lane_dict[lane_id].append(other_lane_id)

    return lane_dict

find_isolated_connections()

Find all isolated strings of connections in the graph. Then Check if any of those lanes would be part of an intersection. Args: None Returns: isolated_connections (list): A list of lists, where each inner list contains the indices of lanes that are part of an isolated connection.

Source code in omega_prime/maposicenterlinesegmentation.py

def find_isolated_connections(self):
    """Find all isolated strings of connections in the graph. Then Check if any of those lanes would be part of an intersection.
    Args:
        None
    Returns:
        isolated_connections (list): A list of lists, where each inner list contains the indices of lanes that are part of an isolated connection.
    """
    G = self.create_non_intersecting_lane_graph()
    new_connections = []
    segment_name_type = type(next(iter(self.lane_segment_dict.values())))

    # Classify each component before constructing ConnectionSegment objects.
    for component in nx.connected_components(G):
        if not component:
            continue

        component_ids = list(component)
        pre = False
        suc = False
        intersection_idxs = set()

        for lane_id in component_ids:
            for successor in self.lane_successors_dict[lane_id]:
                if successor in self.lane_segment_dict and self.lane_segment_dict[successor].segment is not None:
                    intersection_idxs.add(self.lane_segment_dict[successor].segment_idx)
                    suc = True
            for predecessor in self.lane_predecessors_dict[lane_id]:
                if (
                    predecessor in self.lane_segment_dict
                    and self.lane_segment_dict[predecessor].segment is not None
                ):
                    intersection_idxs.add(self.lane_segment_dict[predecessor].segment_idx)
                    pre = True

        if len(intersection_idxs) == 1 and pre and suc:
            # Single bordering intersection on both ends: absorb the component into it.
            inter = self.intersection_dict[list(intersection_idxs)[0]]
            for lane_id in component_ids:
                inter.lanes.append(self.lane_dict[lane_id])
                inter.lane_ids.append(lane_id)
                # Keep lane_segment_dict current so subsequent components see these lanes as already assigned.
                self.lane_segment_dict[lane_id] = segment_name_type(lane_id, inter.idx, inter)
            inter.update_polygon()
        else:
            connection = ConnectionSegment(
                [self.lane_dict[i] for i in component_ids], concave_hull_ratio=self.concave_hull_ratio
            )
            connection.intersection_idxs = intersection_idxs
            new_connections.append(connection)

    isolated_connections = new_connections
    # Create ConnectionSegment for all lanes, that are on multiple intersections:

    isolated_connections = self.combine_isolated_connections(isolated_connections)
    return isolated_connections

get_intersecting_lanes(buffer=None)

Returns a dictionary mapping each lane's lane_id to an array of lane ids it intersects with.

Parameters:

Name	Type	Description	Default
lanes	list	Array of lane objects, each with an idx and centerline attribute.	required

Returns:

Name	Type	Description
dict		A dictionary where keys are lane ids and values are arrays of intersecting lane ids.

Source code in omega_prime/maposicenterlinesegmentation.py

def get_intersecting_lanes(self, buffer: float = None):
    """
    Returns a dictionary mapping each lane's lane_id to an array of lane ids it intersects with.

    Args:
        lanes (list): Array of lane objects, each with an `idx` and `centerline` attribute.

    Returns:
        dict: A dictionary where keys are lane ids and values are arrays of intersecting lane ids.
    """
    if buffer is None:
        buffer = self.lane_buffer

    # Create spatial index directly from centerlines
    lane_centerlines = []
    lane_ids = []

    for lane in self.lanes.values():
        lane_centerlines.append(lane.centerline)
        lane_ids.append(lane.idx.lane_id)

    if not lane_centerlines:
        self.intersecting_lanes_dict = {}
        return {}

    tree = STRtree(lane_centerlines)

    # Pre-compute lane relationships for faster lookup
    successors_set = {lane_id: set(successors) for lane_id, successors in self.lane_successors_dict.items()}
    predecessors_set = {lane_id: set(predecessors) for lane_id, predecessors in self.lane_predecessors_dict.items()}
    parallel_set = {lane_id: set(parallel) for lane_id, parallel in self.parallel_lane_dict.items()}

    intersecting_lanes = {}
    for i, lane in enumerate(self.lanes.values()):
        lane_id = lane.idx.lane_id

        # Query spatial index with buffered geometry
        buffered_centerline = lane.centerline.buffer(buffer)
        candidates = tree.query(buffered_centerline)

        intersecting_lanes[lane_id] = []
        for idx in candidates:
            candidate_id = lane_ids[idx]
            if (
                candidate_id != lane_id
                and candidate_id not in successors_set[lane_id]
                and candidate_id not in predecessors_set[lane_id]
                and candidate_id not in parallel_set[lane_id]
                and buffered_centerline.intersects(lane_centerlines[idx])
            ):
                # Only buffer and test intersection for valid candidates
                intersecting_lanes[lane_id].append(candidate_id)

    self.intersecting_lanes_dict = intersecting_lanes
    return intersecting_lanes

graph_intersection_detection()

Detects intersections in a graph of lanes based on their intersections, successors, and predecessors.

Parameters:

Name	Type	Description	Default
lane_dict	dict	A dictionary where keys are lane indices and values are arrays of intersecting lane indices.	required
lane_successors	dict	A dictionary where keys are lane indices and values are arrays of successor lane indices.	required
lane_predecessors	dict	A dictionary where keys are lane indices and values are arrays of predecessor lane indices.	required

Returns:

Name	Type	Description
list		A list of intersections, where each intersection is a set of lane indices.

Source code in omega_prime/maposicenterlinesegmentation.py

def graph_intersection_detection(self):
    """
    Detects intersections in a graph of lanes based on their intersections, successors, and predecessors.

    Args:
        lane_dict (dict): A dictionary where keys are lane indices and values are arrays of intersecting lane indices.
        lane_successors (dict): A dictionary where keys are lane indices and values are arrays of successor lane indices.
        lane_predecessors (dict): A dictionary where keys are lane indices and values are arrays of predecessor lane indices.

    Returns:
        list: A list of intersections, where each intersection is a set of lane indices.
    """
    # Create a Graph using networkx
    G = nx.Graph()

    # Add nodes and edges to the graph. If a lane has a intersection, add the lanes as nodes and the intersection as an edge
    # Add edges directly (nodes are added automatically)
    for lane_id, intersecting_lanes in self.intersecting_lanes_dict.items():
        G.add_edges_from((lane_id, other_lane) for other_lane in intersecting_lanes)

    intersections = []
    for inter in nx.connected_components(G):
        # Convert lane_ids back to lane objects
        intersection_lanes = [self.lane_dict[i] for i in inter]
        intersection = Intersection(intersection_lanes, concave_hull_ratio=self.concave_hull_ratio)
        intersections.append(intersection)

    self.intersections = intersections
    self.G = G
    return intersections, G

init_intersections()

Initializes the intersections in the map. Args: None Returns: None

Source code in omega_prime/maposicenterlinesegmentation.py

def init_intersections(self):
    """
    Initializes the intersections in the map.
    Args:
        None
    Returns:
        None
    """
    self.create_lane_dict()
    self.get_lane_successors_and_predecessors()
    self.parallel_lane_dict = self.create_parallel_lane_dict()
    self.get_intersecting_lanes()
    self.set_lane_trafficlights()
    self.graph_intersection_detection()
    self.G = add_lanexy_to_graph(self.G, self.lanes)
    self.set_intersection_idx()

    if self.do_combine_intersections:
        self.create_lane_segment_dict()
        self.add_non_intersecting_lanes_to_intersection()
        self.combine_intersections()
        self.set_intersection_idx()

    self.create_intersection_dict()
    self.create_lane_segment_dict()
    self.find_isolated_connections()
    self.create_lane_segment_dict()
    self.check_if_all_lanes_are_on_segment()
    self.update_segment_ids()
    self.create_lane_segment_dict()
    self.update_road_ids()
    self.set_lane_intersection_relation()

intersections_overlap(intersection1, intersection2, buffer=None)

Check if two intersections overlap.

Parameters:

Name	Type	Description	Default
intersection1	Intersection	The first intersection object.	required
intersection2	Intersection	The second intersection object.	required

Returns:

Name	Type	Description
bool		True if the intersections overlap, False otherwise.

Source code in omega_prime/maposicenterlinesegmentation.py

def intersections_overlap(self, intersection1, intersection2, buffer: float = None):
    """
    Check if two intersections overlap.

    Args:
        intersection1 (Intersection): The first intersection object.
        intersection2 (Intersection): The second intersection object.

    Returns:
        bool: True if the intersections overlap, False otherwise.
    """
    if buffer is None:
        buffer = self.intersection_overlap_buffer

    # Use cached buffers if available
    if not hasattr(intersection1, "_buffered_polygon") or intersection1._buffer_value != buffer:
        intersection1._buffered_polygon = intersection1.polygon.buffer(buffer)
        intersection1._buffer_value = buffer

    if not hasattr(intersection2, "_buffered_polygon") or intersection2._buffer_value != buffer:
        intersection2._buffered_polygon = intersection2.polygon.buffer(buffer)
        intersection2._buffer_value = buffer

    return intersection1.polygon.buffer(buffer).intersects(intersection2.polygon.buffer(buffer))

plot(output_plot=None, trajectory=None, plot_lane_ids=False, plot_intersection_polygons=False, plot_connection_polygons=False)

Plots the intersections and saves the plot to the specified output path. A Trajectory can be given to plot it on the map. The Trajectory should be a numpy array of shape (n,3) where each row is (frame, x, y) Args: output_plot (Path): Path to a folder where the plot will be saved. If None, the plot will be shown instead. trajectory (numpy.ndarray): The trajectory to be plotted. If None, no trajectory will be plotted. plot_lane_ids (bool): Whether to plot lane IDs on the map. plot_intersection_polygons (bool): Whether to plot intersection polygons. plot_connection_polygons (bool): Whether to plot connection polygons. Returns: None

Source code in omega_prime/maposicenterlinesegmentation.py

def plot(
    self,
    output_plot: Path = None,
    trajectory=None,
    plot_lane_ids=False,
    plot_intersection_polygons=False,
    plot_connection_polygons=False,
):
    """
    Plots the intersections and saves the plot to the specified output path.
    A Trajectory can be given to plot it on the map. The Trajectory should be a numpy array of shape (n,3) where each row is (frame, x, y)
    Args:
        output_plot (Path): Path to a folder where the plot will be saved. If None, the plot will be shown instead.
        trajectory (numpy.ndarray): The trajectory to be plotted. If None, no trajectory will be plotted.
        plot_lane_ids (bool): Whether to plot lane IDs on the map.
        plot_intersection_polygons (bool): Whether to plot intersection polygons.
        plot_connection_polygons (bool): Whether to plot connection polygons.
    Returns:
        None
    """
    # Plot the map by plotting all the centerlines:
    fig, ax = plt.subplots(1, 1)
    ax.set_aspect(1)

    for lane in self.lanes.values():
        c = "blue"
        if lane.on_intersection:
            c = "green"
        elif lane.is_approaching:
            c = "orange"
        else:
            c = "black"
        ax.plot(*lane.centerline.xy, color=c, alpha=0.3, zorder=-10)

    if plot_lane_ids:
        lane_midpoints = [
            (lane.idx, lane.centerline.interpolate(0.5, normalized=True)) for lane in self.lanes.values()
        ]
        for lane_id, midpoint in lane_midpoints:
            ax.annotate(lane_id, xy=(midpoint.x, midpoint.y), fontsize=2, color="black")

    for inter in self.intersections:
        ax.annotate(inter.idx, xy=inter.get_center_point(), fontsize=2, color="black")

        if plot_intersection_polygons:
            # Plot the polygon into the intersection
            inter.update_polygon()
            ax.plot(*inter.polygon.exterior.xy, color="red", alpha=0.5, zorder=10)

    for combi in self.isolated_connections:
        ax.annotate(combi.idx, xy=combi.get_center_point(), fontsize=2, color="black")
        # Plot the polygon into the intersection
        if plot_connection_polygons:
            combi.update_polygon()
            try:
                ax.plot(*combi.polygon.exterior.xy, color="blue", alpha=0.5, zorder=10)
            except:
                logger.warning(f"Connection {combi.idx} has no polygon")
                pass

    for tl_idx in self.trafficlight:
        position = shapely.Point(
            self.trafficlight[tl_idx].base.position.x, self.trafficlight[tl_idx].base.position.y
        )
        ax.plot(
            position.x,
            position.y,
            marker="o",
            color="red",
            markersize=2,
            label=f"Traffic Light {self.trafficlight[tl_idx].id}",
        )

    # Plot the trajectory if it is given
    if trajectory is not None:
        plt.plot(
            trajectory[:, 1],
            trajectory[:, 2],
            color="yellow",
            alpha=0.8,
            linewidth=3,
            label="Host Vehicle Trajectory",
        )

        # Mark start and end points
        plt.plot(trajectory[0, 1], trajectory[0, 2], "go", markersize=10, label="Start")
        plt.plot(trajectory[-1, 1], trajectory[-1, 2], "ro", markersize=10, label="End")

    ax.set_xlim(*ax.get_xlim())
    ax.set_ylim(*ax.get_ylim())
    plt.title("Map with Intersections")
    plt.xlabel("X Coordinate (m)", fontsize=12)
    plt.ylabel("Y Coordinate (m)", fontsize=12)
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.axis("equal")
    if output_plot is None:
        plt.show()
    else:
        if isinstance(output_plot, Path):
            output_plot.mkdir(parents=True, exist_ok=True)
            plt.savefig(output_plot / "Map_with_Intersection.pdf")
        else:
            isinstance(output_plot, str)
            output_path = Path(output_plot)
            if output_path.is_dir():
                output_path.mkdir(parents=True, exist_ok=True)
                plt.savefig(output_path / "Map_with_Intersection.pdf")
            elif output_path.suffix == ".pdf":
                output_path.parent.mkdir(parents=True, exist_ok=True)
                plt.savefig(output_path)
    plt.close()

plot_intersections(output_plot)

Plots all intersections and saves them to the output path. Args: output_plot (Path): Path to a folder where the plots will be saved. Returns: None

Source code in omega_prime/maposicenterlinesegmentation.py

def plot_intersections(self, output_plot: Path):
    """
    Plots all intersections and saves them to the output path.
    Args:
        output_plot (Path): Path to a folder where the plots will be saved.
    Returns:
        None
    """
    for i, intersection in enumerate(self.intersections):
        intersection.plot(output_plot)
    for i, connection in enumerate(self.isolated_connections):
        connection.plot(output_plot)

set_intersection_idx()

Sets the index for each intersection in the list of intersections. Args: None Returns: None

Source code in omega_prime/maposicenterlinesegmentation.py

def set_intersection_idx(self):
    """
    Sets the index for each intersection in the list of intersections.
    Args:
        None
    Returns:
        None
    """
    for i, intersection in enumerate(self.intersections + self.isolated_connections):
        intersection.idx = i

set_lane_intersection_relation()

Sets the attribute lane.is_approaching true if the lane is connecting to an intersection. Sets the attribute lane.on_intersection true if the lane is part of an intersection.

Source code in omega_prime/maposicenterlinesegmentation.py

def set_lane_intersection_relation(self):
    """
    Sets the attribute lane.is_approaching true if the lane is connecting to an intersection.
    Sets the attribute lane.on_intersection true if the lane is part of an intersection.
    """
    for lane in self.lanes.values():
        self._set_lane_on_intersection(lane, False)
        self._set_lane_is_approaching(lane, False)

    # Process intersection lanes and their predecessors efficiently
    for intersection in self.intersections:
        # Mark intersection lanes
        for lane in intersection.lanes:
            lane_id = self._get_lane_id(lane)
            if lane_id in self.lanes:
                self._set_lane_on_intersection(self.lanes[lane_id], True)

            # Process predecessors for each lane in the intersection
            for predecessor_id in self._get_lane_predecessors(lane):
                if predecessor_id in self.lane_dict:
                    predecessor = self.lane_dict[predecessor_id]
                    if not self._get_lane_on_intersection(predecessor):
                        self._set_lane_is_approaching(predecessor, True)

set_lane_trafficlights()

Sets the traffic lights for each lane of the map.

Source code in omega_prime/maposicenterlinesegmentation.py

def set_lane_trafficlights(self):
    """
    Sets the traffic lights for each lane of the map.
    """
    # Create spatial index for lane centerlines
    lane_centerlines = [lane.centerline for lane in self.lanes.values()]
    lane_objects = list(self.lanes.values())

    tree = STRtree(lane_centerlines)

    for tl_idx in self.trafficlight:
        traffic_light_found = False

        # Create traffic light position point
        tl_point = Point(self.trafficlight[tl_idx].base.position.x, self.trafficlight[tl_idx].base.position.y)

        # Use spatial index to find candidate lanes
        candidates = tree.nearest(tl_point)

        if candidates:
            lane = lane_objects[candidates]
            lane.traffic_light = self.trafficlight[tl_idx]
            traffic_light_found = True

        if not traffic_light_found:
            logger.warning(f"Traffic light {self.trafficlight[tl_idx].id} not found in any lane")

trajectory_segment_detection(trajectory)

Splits a trajectory into segments based on the lane it is located on

Parameters:

Name	Type	Description	Default
trajectory	ndarray	A NumPy array of shape (n, 3) representing the trajectory, where each row is a (frame, x, y) coordinate.	required

Returns:

Name	Type	Description
list		A list of tuples, where each tuple contains a segment of the trajectory and the segment it intersects with.

Source code in omega_prime/maposicenterlinesegmentation.py

def trajectory_segment_detection(self, trajectory):
    """
    Splits a trajectory into segments based on the lane it is located on

    Args:
        trajectory (np.ndarray): A NumPy array of shape (n, 3) representing the trajectory, where each row is a (frame, x, y) coordinate.

    Returns:
        list: A list of tuples, where each tuple contains a segment of the trajectory and the segment it intersects with.
    """
    segments = []
    current_segment = []
    xy = trajectory[:, 1:3]  # Extract x and y coordinates
    sts = self.locator.xys2sts(xy)
    lane_ids = sts["roadlane_id"].to_numpy()
    segment_idx = [self.lane_segment_dict[lane_id.lane_id].segment.idx for lane_id in lane_ids]

    trajectory = np.column_stack((trajectory[:, 0], trajectory[:, 1], trajectory[:, 2], lane_ids, segment_idx))

    # Create spatial index for intersection polygons
    intersection_polygons = []
    intersection_ids = []
    buffer = 5

    for segment in self.segments:
        if segment.type == MapSegmentType.JUNCTION and hasattr(segment, "polygon"):
            intersection_polygons.append(segment.polygon.buffer(buffer))
            intersection_ids.append(segment.idx)

    if intersection_polygons:
        # Use spatial index for efficient intersection queries
        tree = STRtree(intersection_polygons)

        # Process points in batches for better performance
        for i, (frame, x, y, _, _) in enumerate(trajectory):
            point = Point(x, y)

            # Query spatial index instead of checking all polygons
            candidates = tree.query(point)

            for idx in candidates:
                if intersection_polygons[idx].contains(point):
                    trajectory[i, 4] = intersection_ids[idx]
                    break

    # Rest of the method for creating segments
    prev_seg_id = -1
    for i, (frame, x, y, _, segment_idx) in enumerate(trajectory):
        if prev_seg_id == segment_idx:
            current_segment.append((frame, x, y))
        else:
            if current_segment:
                segments.append((np.array(current_segment), self.segments[prev_seg_id]))
            current_segment = [(frame, x, y)]
            prev_seg_id = segment_idx

    if current_segment:
        segments.append((np.array(current_segment), self.segments[prev_seg_id]))

    return segments

update_road_ids()

Updates the road_ids of the lane to the segment ID

Source code in omega_prime/maposicenterlinesegmentation.py

def update_road_ids(self):
    """
    Updates the road_ids of the lane to the segment ID
    """
    updates_needed = []
    old_to_new_mapping = {}

    # First pass: identify what needs to be updated
    for lane_idx, lane in self.lanes.items():
        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:
            new_road_id = self.lane_segment_dict[lane_id].segment.idx
            if lane.idx.road_id != new_road_id:
                new_idx = lane.idx._replace(road_id=new_road_id)
                updates_needed.append((lane_idx, lane, new_idx))
                old_to_new_mapping[lane_idx] = new_idx

    # Second pass: apply updates efficiently
    for old_idx, lane, new_idx in updates_needed:
        # Update the lane object in place
        lane.idx = new_idx

        # Only modify dictionary if the key actually changed
        if old_idx != new_idx:
            self.lanes[new_idx] = lane
            del self.lanes[old_idx]

    # Third pass: update all predecessor and successor references
    for lane in self.lanes.values():
        # Update predecessor references
        updated_predecessors = []
        for pred_id in lane.predecessor_ids:
            if pred_id in old_to_new_mapping:
                updated_predecessors.append(old_to_new_mapping[pred_id])
            else:
                updated_predecessors.append(pred_id)
        lane.predecessor_ids = updated_predecessors

        # Update successor references
        updated_successors = []
        for succ_id in lane.successor_ids:
            if succ_id in old_to_new_mapping:
                updated_successors.append(old_to_new_mapping[succ_id])
            else:
                updated_successors.append(succ_id)
        lane.successor_ids = updated_successors

    # Fourth pass: update internal dictionaries that track relationships
    self.lane_dict = {lane.idx.lane_id: lane for lane in self.lanes.values()}
    self.get_lane_successors_and_predecessors()

update_segment_ids()

Updates the segment IDs of the map segmentation

Source code in omega_prime/maposicenterlinesegmentation.py

def update_segment_ids(self):
    "Updates the segment IDs of the map segmentation"
    self.segments = self.intersections + self.isolated_connections
    for i, segment in enumerate(self.segments):
        segment.idx = i
        segment.set_trafficlight()

SegmentOsiCenterline

Bases: Segment

Segment implementation for OSI centerline-based maps.

Source code in omega_prime/maposicenterlinesegmentation.py

class SegmentOsiCenterline(Segment):
    """Segment implementation for OSI centerline-based maps."""

    def _get_lane_id(self, lane):
        return lane.idx.lane_id

    def _get_lane_geometry(self, lane) -> shapely.LineString:
        return lane.centerline

    def set_trafficlight(self):
        trafficlight = []
        for lane in self.lanes:
            if hasattr(lane, "trafficlight") and lane.trafficlight:
                trafficlight.append(lane.trafficlight)
        self.trafficlights = trafficlight

add_lanexy_to_graph(G, lanes)

Adds lane coordinates to the graph as node attributes.

Parameters:

Name	Type	Description	Default
G	Graph	The graph to which lane coordinates will be added.	required
lanes	dict	A dictionary of lane objects.	required

Returns:

Type	Description
	networkx.Graph: The updated graph with lane coordinates as node attributes.

Source code in omega_prime/maposicenterlinesegmentation.py

def add_lanexy_to_graph(G: nx.Graph, lanes):
    """
    Adds lane coordinates to the graph as node attributes.

    Args:
        G (networkx.Graph): The graph to which lane coordinates will be added.
        lanes (dict): A dictionary of lane objects.

    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

plot_graph(G, output)

Plots the graph with lane coordinates.

Parameters:

Name	Type	Description	Default
G	Graph	The graph to be plotted.	required
Path	str or Path	The file path to save the plot.	required

Source code in omega_prime/maposicenterlinesegmentation.py

def plot_graph(G: nx.Graph, output: 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