328 lines
10 KiB
Python
328 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate a compact road-class grid for the RED4ext GPS edge-cost shim."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
|
|
|
|
FLAGS = {
|
|
"Pavement": 0x0008,
|
|
"Road": 0x0010,
|
|
"GPSOnly": 0x0200,
|
|
"Highway": 0x4000,
|
|
}
|
|
|
|
CATEGORY_CODES = {
|
|
"unknown": ".",
|
|
"pavement": "P",
|
|
"gpsonly": "G",
|
|
"road": "R",
|
|
"highway": "H",
|
|
}
|
|
|
|
CATEGORY_PRIORITY = {
|
|
".": 0,
|
|
"P": 1,
|
|
"G": 2,
|
|
"R": 3,
|
|
"H": 4,
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class GridSpec:
|
|
min_x: float
|
|
min_y: float
|
|
width: int
|
|
height: int
|
|
cell_size: float
|
|
|
|
|
|
def flag_value(value: Any) -> int:
|
|
if isinstance(value, int):
|
|
return value
|
|
if isinstance(value, str):
|
|
return int(value)
|
|
if isinstance(value, dict):
|
|
for key in ("Value", "$value", "value"):
|
|
if key in value:
|
|
return flag_value(value[key])
|
|
return int(value or 0)
|
|
|
|
|
|
def lane_category(flags: int) -> str:
|
|
if flags & FLAGS["Highway"]:
|
|
return "highway"
|
|
if flags & FLAGS["GPSOnly"]:
|
|
return "gpsonly"
|
|
if flags & FLAGS["Road"]:
|
|
return "road"
|
|
if flags & FLAGS["Pavement"]:
|
|
return "pavement"
|
|
return "unknown"
|
|
|
|
|
|
def load_lanes(path: Path) -> list[dict[str, Any]]:
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
return data["Data"]["RootChunk"]["data"]["lanes"]
|
|
|
|
|
|
def load_lane_polygons(path: Path) -> list[list[tuple[float, float]]]:
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
rows = data["Data"]["RootChunk"]["data"]
|
|
polygons: list[list[tuple[float, float]]] = []
|
|
for row in rows:
|
|
polygon = row.get("value", {}).get("polygon") or []
|
|
polygons.append([(float(point["X"]), float(point["Y"])) for point in polygon])
|
|
return polygons
|
|
|
|
|
|
def point_in_polygon(x: float, y: float, polygon: list[tuple[float, float]]) -> bool:
|
|
inside = False
|
|
count = len(polygon)
|
|
if count < 3:
|
|
return False
|
|
|
|
previous_x, previous_y = polygon[-1]
|
|
for current_x, current_y in polygon:
|
|
if (current_y > y) != (previous_y > y):
|
|
edge_x = (previous_x - current_x) * (y - current_y) / (previous_y - current_y) + current_x
|
|
if x < edge_x:
|
|
inside = not inside
|
|
previous_x, previous_y = current_x, current_y
|
|
|
|
return inside
|
|
|
|
|
|
def orientation(ax: float, ay: float, bx: float, by: float, cx: float, cy: float) -> float:
|
|
return (by - ay) * (cx - bx) - (bx - ax) * (cy - by)
|
|
|
|
|
|
def on_segment(ax: float, ay: float, bx: float, by: float, cx: float, cy: float) -> bool:
|
|
return min(ax, cx) <= bx <= max(ax, cx) and min(ay, cy) <= by <= max(ay, cy)
|
|
|
|
|
|
def segments_intersect(
|
|
ax: float,
|
|
ay: float,
|
|
bx: float,
|
|
by: float,
|
|
cx: float,
|
|
cy: float,
|
|
dx: float,
|
|
dy: float,
|
|
) -> bool:
|
|
o1 = orientation(ax, ay, bx, by, cx, cy)
|
|
o2 = orientation(ax, ay, bx, by, dx, dy)
|
|
o3 = orientation(cx, cy, dx, dy, ax, ay)
|
|
o4 = orientation(cx, cy, dx, dy, bx, by)
|
|
|
|
if (o1 > 0) != (o2 > 0) and (o3 > 0) != (o4 > 0):
|
|
return True
|
|
|
|
eps = 1e-5
|
|
if abs(o1) <= eps and on_segment(ax, ay, cx, cy, bx, by):
|
|
return True
|
|
if abs(o2) <= eps and on_segment(ax, ay, dx, dy, bx, by):
|
|
return True
|
|
if abs(o3) <= eps and on_segment(cx, cy, ax, ay, dx, dy):
|
|
return True
|
|
if abs(o4) <= eps and on_segment(cx, cy, bx, by, dx, dy):
|
|
return True
|
|
return False
|
|
|
|
|
|
def polygon_intersects_cell(
|
|
polygon: list[tuple[float, float]],
|
|
min_x: float,
|
|
min_y: float,
|
|
max_x: float,
|
|
max_y: float,
|
|
) -> bool:
|
|
center_x = (min_x + max_x) * 0.5
|
|
center_y = (min_y + max_y) * 0.5
|
|
if point_in_polygon(center_x, center_y, polygon):
|
|
return True
|
|
|
|
corners = ((min_x, min_y), (max_x, min_y), (max_x, max_y), (min_x, max_y))
|
|
if any(point_in_polygon(x, y, polygon) for x, y in corners):
|
|
return True
|
|
|
|
if any(min_x <= x <= max_x and min_y <= y <= max_y for x, y in polygon):
|
|
return True
|
|
|
|
cell_edges = (
|
|
(min_x, min_y, max_x, min_y),
|
|
(max_x, min_y, max_x, max_y),
|
|
(max_x, max_y, min_x, max_y),
|
|
(min_x, max_y, min_x, min_y),
|
|
)
|
|
previous_x, previous_y = polygon[-1]
|
|
for current_x, current_y in polygon:
|
|
for edge in cell_edges:
|
|
if segments_intersect(previous_x, previous_y, current_x, current_y, *edge):
|
|
return True
|
|
previous_x, previous_y = current_x, current_y
|
|
|
|
return False
|
|
|
|
|
|
def grid_index(spec: GridSpec, x: float, y: float) -> tuple[int, int]:
|
|
column = math.floor((x - spec.min_x) / spec.cell_size)
|
|
row = math.floor((y - spec.min_y) / spec.cell_size)
|
|
return int(column), int(row)
|
|
|
|
|
|
def mark_cell(grid: list[list[str]], column: int, row: int, code: str) -> bool:
|
|
if row < 0 or row >= len(grid) or column < 0 or column >= len(grid[row]):
|
|
return False
|
|
if CATEGORY_PRIORITY[code] <= CATEGORY_PRIORITY[grid[row][column]]:
|
|
return False
|
|
grid[row][column] = code
|
|
return True
|
|
|
|
|
|
def iter_cells_for_polygon(spec: GridSpec, polygon: list[tuple[float, float]]) -> Iterable[tuple[int, int]]:
|
|
xs = [point[0] for point in polygon]
|
|
ys = [point[1] for point in polygon]
|
|
min_col, min_row = grid_index(spec, min(xs), min(ys))
|
|
max_col, max_row = grid_index(spec, max(xs), max(ys))
|
|
|
|
for row in range(max(0, min_row), min(spec.height - 1, max_row) + 1):
|
|
for column in range(max(0, min_col), min(spec.width - 1, max_col) + 1):
|
|
cell_min_x = spec.min_x + column * spec.cell_size
|
|
cell_min_y = spec.min_y + row * spec.cell_size
|
|
if polygon_intersects_cell(
|
|
polygon,
|
|
cell_min_x,
|
|
cell_min_y,
|
|
cell_min_x + spec.cell_size,
|
|
cell_min_y + spec.cell_size,
|
|
):
|
|
yield column, row
|
|
|
|
|
|
def build_grid(
|
|
lanes: list[dict[str, Any]],
|
|
polygons: list[list[tuple[float, float]]],
|
|
cell_size: float,
|
|
inflate_cells: int,
|
|
) -> tuple[GridSpec, list[list[str]], dict[str, int]]:
|
|
points = [point for polygon in polygons for point in polygon]
|
|
if not points:
|
|
raise ValueError("no polygon points found")
|
|
|
|
min_x = math.floor(min(point[0] for point in points) / cell_size) * cell_size
|
|
min_y = math.floor(min(point[1] for point in points) / cell_size) * cell_size
|
|
max_x = math.ceil(max(point[0] for point in points) / cell_size) * cell_size
|
|
max_y = math.ceil(max(point[1] for point in points) / cell_size) * cell_size
|
|
spec = GridSpec(
|
|
min_x=min_x,
|
|
min_y=min_y,
|
|
width=int(round((max_x - min_x) / cell_size)),
|
|
height=int(round((max_y - min_y) / cell_size)),
|
|
cell_size=cell_size,
|
|
)
|
|
|
|
grid = [[CATEGORY_CODES["unknown"] for _ in range(spec.width)] for _ in range(spec.height)]
|
|
for index, (lane, polygon) in enumerate(zip(lanes, polygons, strict=True)):
|
|
if len(polygon) < 3:
|
|
continue
|
|
code = CATEGORY_CODES[lane_category(flag_value(lane.get("flags", 0)))]
|
|
if code == CATEGORY_CODES["unknown"]:
|
|
continue
|
|
|
|
changed_cells: set[tuple[int, int]] = set()
|
|
for column, row in iter_cells_for_polygon(spec, polygon):
|
|
for dy in range(-inflate_cells, inflate_cells + 1):
|
|
for dx in range(-inflate_cells, inflate_cells + 1):
|
|
changed_cells.add((column + dx, row + dy))
|
|
|
|
for column, row in changed_cells:
|
|
mark_cell(grid, column, row, code)
|
|
|
|
counts = {code: sum(row.count(code) for row in grid) for code in CATEGORY_PRIORITY}
|
|
return spec, grid, counts
|
|
|
|
|
|
def write_header(path: Path, spec: GridSpec, grid: list[list[str]], counts: dict[str, int]) -> None:
|
|
lines = [
|
|
"#pragma once",
|
|
"",
|
|
"#include <array>",
|
|
"#include <cstdint>",
|
|
"#include <string_view>",
|
|
"",
|
|
"namespace EdgeWeightGPS::Generated",
|
|
"{",
|
|
f"constexpr float kSpatialRoadGridMinX = {spec.min_x:.1f}f;",
|
|
f"constexpr float kSpatialRoadGridMinY = {spec.min_y:.1f}f;",
|
|
f"constexpr float kSpatialRoadGridCellSize = {spec.cell_size:.1f}f;",
|
|
f"constexpr uint32_t kSpatialRoadGridWidth = {spec.width};",
|
|
f"constexpr uint32_t kSpatialRoadGridHeight = {spec.height};",
|
|
"",
|
|
"// Cell codes: H=highway, R=road, G=GPS-only, P=pavement, .=unknown.",
|
|
"// Generated by tools/generate_spatial_edge_grid.py from all.traffic_persistent and all.lane_polygons.",
|
|
f"// Counts: H={counts['H']} R={counts['R']} G={counts['G']} P={counts['P']} unknown={counts['.']}.",
|
|
"constexpr std::array<std::string_view, kSpatialRoadGridHeight> kSpatialRoadGridRows = {{",
|
|
]
|
|
|
|
for row in grid:
|
|
lines.append(f' "{("").join(row)}",')
|
|
|
|
lines.extend(
|
|
[
|
|
"}};",
|
|
"",
|
|
"} // namespace EdgeWeightGPS::Generated",
|
|
"",
|
|
]
|
|
)
|
|
path.write_text("\n".join(lines), encoding="utf-8")
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--traffic-json", type=Path, default=Path("work/raw-segment-json/all.traffic_persistent.json"))
|
|
parser.add_argument(
|
|
"--lane-polygons-json",
|
|
type=Path,
|
|
default=Path("work/traffic-companions/json/all.lane_polygons.json"),
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
type=Path,
|
|
default=Path("red4ext/EdgeWeightGPS/src/GeneratedSpatialRoadGrid.hpp"),
|
|
)
|
|
parser.add_argument("--cell-size", type=float, default=16.0)
|
|
parser.add_argument("--inflate-cells", type=int, default=0)
|
|
args = parser.parse_args()
|
|
|
|
lanes = load_lanes(args.traffic_json)
|
|
polygons = load_lane_polygons(args.lane_polygons_json)
|
|
if len(lanes) != len(polygons):
|
|
raise ValueError(f"lane/polygon count mismatch: {len(lanes)} != {len(polygons)}")
|
|
|
|
spec, grid, counts = build_grid(lanes, polygons, args.cell_size, args.inflate_cells)
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
write_header(args.output, spec, grid, counts)
|
|
|
|
total = spec.width * spec.height
|
|
covered = total - sum(row.count(".") for row in grid)
|
|
print(
|
|
f"wrote {args.output} cells={total} covered={covered} "
|
|
f"size={spec.width}x{spec.height} origin=({spec.min_x},{spec.min_y}) cell={spec.cell_size}"
|
|
)
|
|
print("counts " + " ".join(f"{code}={counts[code]}" for code in ("H", "R", "G", "P", ".")))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|