133 lines
3.7 KiB
Python
Executable File
133 lines
3.7 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Patch traffic lane connection probabilities toward highway routing."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
FLAGS = {
|
|
"Road": 16,
|
|
"Pavement": 8,
|
|
"TrafficDisabled": 128,
|
|
"CrossWalk": 256,
|
|
"GPSOnly": 512,
|
|
"Blockade": 2048,
|
|
"Highway": 16384,
|
|
}
|
|
|
|
|
|
def load_json(path: Path) -> Any:
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def flag_value(flags: Any) -> int:
|
|
return int(flags or 0)
|
|
|
|
|
|
def has(flags: int, name: str) -> bool:
|
|
return bool(flags & FLAGS[name])
|
|
|
|
|
|
def category(flags: int) -> str:
|
|
if has(flags, "Highway"):
|
|
return "highway"
|
|
if has(flags, "GPSOnly"):
|
|
return "gpsonly"
|
|
if has(flags, "Road"):
|
|
return "road"
|
|
if has(flags, "Pavement"):
|
|
return "pavement"
|
|
return "other"
|
|
|
|
|
|
def write_json(path: Path, data: Any) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(data, indent=2) + "\n", encoding="utf-8")
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("traffic_persistent_json", type=Path)
|
|
parser.add_argument("lane_connections_json", type=Path)
|
|
parser.add_argument("out_json", type=Path)
|
|
parser.add_argument("--high", type=int, default=255)
|
|
parser.add_argument("--low", type=int, default=1)
|
|
args = parser.parse_args()
|
|
|
|
traffic = load_json(args.traffic_persistent_json)
|
|
lanes = traffic["Data"]["RootChunk"]["data"]["lanes"]
|
|
lane_categories = [category(flag_value(lane.get("flags", 0))) for lane in lanes]
|
|
|
|
connections = load_json(args.lane_connections_json)
|
|
rows = connections["Data"]["RootChunk"]["data"]
|
|
|
|
changed = 0
|
|
high_set = 0
|
|
low_set = 0
|
|
touched_rows = 0
|
|
|
|
for row in rows:
|
|
source_index = row["index"]
|
|
source_category = lane_categories[source_index]
|
|
outlanes = row["value"]["outlanes"]
|
|
if not outlanes:
|
|
continue
|
|
|
|
target_categories = [
|
|
lane_categories[out["laneIndex"]]
|
|
for out in outlanes
|
|
if 0 <= out["laneIndex"] < len(lane_categories)
|
|
]
|
|
has_highway_target = "highway" in target_categories
|
|
has_non_highway_target = any(cat != "highway" for cat in target_categories)
|
|
row_changed = False
|
|
|
|
for out in outlanes:
|
|
target_category = lane_categories[out["laneIndex"]]
|
|
old = out["exitProbabilityCompressed"]
|
|
new = old
|
|
|
|
if target_category == "highway" and source_category in {"road", "gpsonly", "highway"}:
|
|
new = args.high
|
|
elif (
|
|
source_category == "highway"
|
|
and has_highway_target
|
|
and target_category != "highway"
|
|
):
|
|
new = args.low
|
|
elif (
|
|
source_category in {"road", "gpsonly"}
|
|
and has_highway_target
|
|
and has_non_highway_target
|
|
and target_category != "highway"
|
|
):
|
|
new = args.low
|
|
|
|
if new != old:
|
|
out["exitProbabilityCompressed"] = new
|
|
changed += 1
|
|
row_changed = True
|
|
if new == args.high:
|
|
high_set += 1
|
|
elif new == args.low:
|
|
low_set += 1
|
|
|
|
if row_changed:
|
|
touched_rows += 1
|
|
|
|
write_json(args.out_json, connections)
|
|
print(f"rows seen: {len(rows)}")
|
|
print(f"rows touched: {touched_rows}")
|
|
print(f"edges changed: {changed}")
|
|
print(f"edges set high: {high_set}")
|
|
print(f"edges set low: {low_set}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|