|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import xml.etree.ElementTree as ET |
| 4 | +from collections import defaultdict |
| 5 | +from typing import Final |
| 6 | + |
| 7 | +from attrs import define, field |
| 8 | + |
| 9 | +from t4_devkit.typing import Vector3 |
| 10 | + |
| 11 | +__all__ = ["LaneletParser"] |
| 12 | + |
| 13 | + |
| 14 | +@define |
| 15 | +class Node: |
| 16 | + """Represents an OSM node (point) with coordinates and attributes.""" |
| 17 | + |
| 18 | + id: str |
| 19 | + lat: float |
| 20 | + lon: float |
| 21 | + local_x: float | None = None |
| 22 | + local_y: float | None = None |
| 23 | + ele: float | None = None |
| 24 | + tags: dict[str, str] = field(factory=dict) |
| 25 | + |
| 26 | + |
| 27 | +@define |
| 28 | +class Way: |
| 29 | + """Represents an OSM way (line/polyline) with node references and attributes.""" |
| 30 | + |
| 31 | + id: str |
| 32 | + node_refs: list[str] = field(factory=list) |
| 33 | + tags: dict[str, str] = field(factory=dict) |
| 34 | + |
| 35 | + |
| 36 | +@define |
| 37 | +class Relation: |
| 38 | + """Represents an OSM relation with member references and attributes.""" |
| 39 | + |
| 40 | + id: str |
| 41 | + members: list[Member] = field(factory=list) |
| 42 | + tags: dict[str, str] = field(factory=dict) |
| 43 | + |
| 44 | + |
| 45 | +@define |
| 46 | +class Member: |
| 47 | + """Represents an OSM relation member.""" |
| 48 | + |
| 49 | + type: str |
| 50 | + ref: str |
| 51 | + role: str |
| 52 | + |
| 53 | + |
| 54 | +class LaneletParser: |
| 55 | + """Parses an OSM XML file into a dictionary of nodes, ways, and relations.""" |
| 56 | + |
| 57 | + elevation_scale: float = 1.0 |
| 58 | + default_elevation: float = 0.0 |
| 59 | + |
| 60 | + def __init__(self, filepath: str, *, verbose: bool = False): |
| 61 | + """Initializes the parser with the given file path. |
| 62 | +
|
| 63 | + Args: |
| 64 | + filepath (str): The path to the OSM XML file to parse. |
| 65 | + verbose (bool, optional): Whether to print basic statistics about the parsed OSM data. |
| 66 | + """ |
| 67 | + |
| 68 | + tree = ET.parse(filepath) |
| 69 | + root = tree.getroot() |
| 70 | + |
| 71 | + self._nodes = _parse_nodes(root) |
| 72 | + self._ways = _parse_ways(root) |
| 73 | + self._relations = _parse_relations(root) |
| 74 | + |
| 75 | + if verbose: |
| 76 | + self._print_statistics() |
| 77 | + |
| 78 | + def _print_statistics(self) -> None: |
| 79 | + """Print basic statistics about the parsed OSM data.""" |
| 80 | + num_lines: Final[int] = 50 |
| 81 | + |
| 82 | + print("\n" + "=" * num_lines) |
| 83 | + print("OSM MAP STATISTICS") |
| 84 | + print("=" * num_lines) |
| 85 | + print(f"Nodes: {len(self.nodes)}") |
| 86 | + print(f"Ways: {len(self.ways)}") |
| 87 | + print(f"Relations: {len(self.relations)}") |
| 88 | + |
| 89 | + # Analyze way types |
| 90 | + way_types = defaultdict(int) |
| 91 | + for way in self.ways.values(): |
| 92 | + way_type = way.tags.get("type", "unknown") |
| 93 | + subtype = way.tags.get("subtype", "") |
| 94 | + key = f"{way_type}:{subtype}" if subtype else way_type |
| 95 | + way_types[key] += 1 |
| 96 | + |
| 97 | + print("\nWay Types:") |
| 98 | + for way_type, count in sorted(way_types.items()): |
| 99 | + print(f" {way_type}: {count}") |
| 100 | + |
| 101 | + # Analyze relation types |
| 102 | + relation_types = defaultdict(int) |
| 103 | + for relation in self.relations.values(): |
| 104 | + rel_type = relation.tags.get("type", "unknown") |
| 105 | + subtype = relation.tags.get("subtype", "") |
| 106 | + key = f"{rel_type}:{subtype}" if subtype else rel_type |
| 107 | + relation_types[key] += 1 |
| 108 | + |
| 109 | + print("\nRelation Types:") |
| 110 | + for rel_type, count in sorted(relation_types.items()): |
| 111 | + print(f" {rel_type}: {count}") |
| 112 | + |
| 113 | + # Coordinate system info |
| 114 | + local_coord_nodes = sum( |
| 115 | + 1 |
| 116 | + for node in self.nodes.values() |
| 117 | + if node.local_x is not None and node.local_y is not None |
| 118 | + ) |
| 119 | + print(f"\nNodes with local coordinates: {local_coord_nodes}/{len(self.nodes)}") |
| 120 | + print("=" * num_lines) |
| 121 | + |
| 122 | + @property |
| 123 | + def nodes(self) -> dict[str, Node]: |
| 124 | + return self._nodes |
| 125 | + |
| 126 | + @property |
| 127 | + def ways(self) -> dict[str, Way]: |
| 128 | + return self._ways |
| 129 | + |
| 130 | + @property |
| 131 | + def relations(self) -> dict[str, Relation]: |
| 132 | + return self._relations |
| 133 | + |
| 134 | + def coordinates(self, node: Node, *, as_geographic: bool = False) -> Vector3: |
| 135 | + """Return coordinates of a node, preferring local coordinates if available. |
| 136 | +
|
| 137 | + Args: |
| 138 | + node (Node): The node to get coordinates for. |
| 139 | + as_geographic (bool): Whether to return coordinates in geographic (lat, lon, elevation) format. |
| 140 | +
|
| 141 | + Returns: |
| 142 | + A Vector3 coordinate for the node. |
| 143 | + """ |
| 144 | + if node.local_x is not None and node.local_y is not None and not as_geographic: |
| 145 | + x, y = node.local_x, node.local_y |
| 146 | + else: |
| 147 | + # Convert lat/lon to a simple projection (not accurate for large areas) |
| 148 | + x, y = node.lat, node.lon |
| 149 | + |
| 150 | + z = node.ele * self.elevation_scale if node.ele is not None else self.default_elevation |
| 151 | + |
| 152 | + return Vector3(x, y, z) |
| 153 | + |
| 154 | + def way_coordinates(self, way: Way, *, as_geographic: bool = False) -> list[Vector3]: |
| 155 | + """Return coordinates of a way, preferring local coordinates if available. |
| 156 | +
|
| 157 | + Args: |
| 158 | + way (Way): The way to get coordinates for. |
| 159 | + as_geographic: Whether to return coordinates in geographic (lat, lon, elevation) format. |
| 160 | +
|
| 161 | + Returns: |
| 162 | + A list of Vector3 coordinates for the way. |
| 163 | + """ |
| 164 | + return [ |
| 165 | + self.coordinates(self.nodes[node_ref], as_geographic=as_geographic) |
| 166 | + for node_ref in way.node_refs |
| 167 | + if node_ref in self.nodes |
| 168 | + ] |
| 169 | + |
| 170 | + |
| 171 | +def _parse_nodes(root: ET.Element) -> dict[str, Node]: |
| 172 | + output: dict[str, Node] = {} |
| 173 | + for node_elem in root.findall("node"): |
| 174 | + node_id = node_elem.get("id") |
| 175 | + node_lat = float(node_elem.get("lat")) |
| 176 | + node_lon = float(node_elem.get("lon")) |
| 177 | + |
| 178 | + tags: dict[str, str] = {} |
| 179 | + local_x = None |
| 180 | + local_y = None |
| 181 | + ele = None |
| 182 | + for tag_elem in node_elem.findall("tag"): |
| 183 | + key = tag_elem.get("k") |
| 184 | + value = tag_elem.get("v") |
| 185 | + tags[key] = value |
| 186 | + |
| 187 | + # extract local coordinates if available |
| 188 | + if key == "local_x": |
| 189 | + local_x = float(value) |
| 190 | + elif key == "local_y": |
| 191 | + local_y = float(value) |
| 192 | + elif key == "ele": |
| 193 | + ele = float(value) |
| 194 | + |
| 195 | + output[node_id] = Node( |
| 196 | + id=node_id, |
| 197 | + lat=node_lat, |
| 198 | + lon=node_lon, |
| 199 | + tags=tags, |
| 200 | + local_x=local_x, |
| 201 | + local_y=local_y, |
| 202 | + ele=ele, |
| 203 | + ) |
| 204 | + |
| 205 | + return output |
| 206 | + |
| 207 | + |
| 208 | +def _parse_ways(root: ET.Element) -> dict[str, Way]: |
| 209 | + output: dict[str, Way] = {} |
| 210 | + for way_elem in root.findall("way"): |
| 211 | + way_id = way_elem.get("id") |
| 212 | + node_refs = [nd.get("ref") for nd in way_elem.findall("nd")] |
| 213 | + |
| 214 | + tags: dict[str, str] = { |
| 215 | + tag_elem.get("k"): tag_elem.get("v") for tag_elem in way_elem.findall("tag") |
| 216 | + } |
| 217 | + |
| 218 | + output[way_id] = Way(id=way_id, node_refs=node_refs, tags=tags) |
| 219 | + |
| 220 | + return output |
| 221 | + |
| 222 | + |
| 223 | +def _parse_relations(root: ET.Element) -> dict[str, Relation]: |
| 224 | + output: dict[str, Relation] = {} |
| 225 | + for relation_elem in root.findall("relation"): |
| 226 | + relation_id = relation_elem.get("id") |
| 227 | + |
| 228 | + members = [ |
| 229 | + Member( |
| 230 | + type=member_elem.get("type"), |
| 231 | + ref=member_elem.get("ref"), |
| 232 | + role=member_elem.get("role", ""), |
| 233 | + ) |
| 234 | + for member_elem in relation_elem.findall("member") |
| 235 | + ] |
| 236 | + |
| 237 | + tags: dict[str, str] = { |
| 238 | + tag_elem.get("k"): tag_elem.get("v") for tag_elem in relation_elem.findall("tag") |
| 239 | + } |
| 240 | + |
| 241 | + output[relation_id] = Relation(id=relation_id, members=members, tags=tags) |
| 242 | + |
| 243 | + return output |
0 commit comments