Skip to content

Repository files navigation

sldm-gnn

AI-based trigger for the S-LDM (DriveX-devs/S-LDM) connected-vehicle service, developed for the MSc thesis "AI events prediction for centralized Server Local Dynamic Map connected vehicle service".

The S-LDM already exposes an asynchronous-trigger hook for critical driving events, but currently the trigger is only a demonstrative version based on a simple hardcoded rule-based mechanism. This repo provides a primitive implementation of the AI core that is expected to power the next version of the trigger: a Graph Neural Network (GNN) that ingests data from the S-LDM and predicts the occurrence of one of four critical events (encoded as a bitmask in the dataset's MLBEncoded column, see src/labels.py):

Index Label Bit value
0 LANE_CHANGE 1
1 OVERTAKE 2
2 TURN 4
3 COLLISION 8

Each run targets a single label, so the model is trained as a binary classifier.

Installation

The project is managed with uv, which fetches the required Python (3.13, per .python-version and requires-python) and builds an isolated project venv (.venv/) to achieve isolation and reproducibility, and preventing conflicts. No system Python or global installs needed, avoiding interference with the host's interpreter and libraries.

To install uv:

curl -LsSf https://astral.sh/uv/install.sh | sh

Then, once cloned the repository, the dependencies can be easiliy installed with the provided script, that install the correct versions of the libaries automatically detecting the available hardware:

./install_dependencies.sh   # auto-detects CUDA; uses uv sync --extra cuda|cpu

Then, to automatically run python scripts with the correct venv, the user is encouraged to run each script with uv run <script> rather than python3 <script>.

Model & code layout

GruSage (src/models/grusage.py) performs the aggregation sequentially: first, the GRU aggregates on the temporal dimension the dynamic features contained in each vehicle/node; static features of each vehicle/node does not need temporal aggregation at all, the only additional module for them is the embedding for the categorical station type; then, optionally, each vehicle can receive neighborhood-contextual map features, which are eventually pre-processed with an encoder and buffereed and then selected for each vehicle with a top-k closest spatial attention mechanism related to last-frame position; finally, the SAGE layers implement graph convolution, and the subsequent graph pooling module aggregates over the vehicle graph returning one array (on the feature dimension) for each graph.

build.py, main.py, test.py, compare_predictions_gt.py, rcv.py  # CLI entrypoints
install_dependencies.sh                                          # uv installer (CPU/CUDA auto-detect)
src/
  labels.py         # LabelsEnum
  gbuilder.py       # graph building: classes, functions and utilities
  dataset.py        # graph dataset
  models/           # GruSage + map encoder/attention
  transforms.py     # data augmentation and load transforms
  metrics.py        # PackMetrics (per-pack) + EventMetrics (event-level clusters)
  utils.py          # miscellaneous: training loop, MetaData, ParamSweepContext, ...

Data layout

Each directory passed to build.py should have shape as in the following:

<data_path>/
  vmap.parquet                # lane map (consumed by MapBuilder)
  train/ valid/ test/         # each with:
    packs.parquet             # (PackId, FrameId, VehicleId, X, Y, Speed, Angle, ...)
    labels.parquet            # (PackId, MLBEncoded)
    vinfo.parquet             # (VehicleId, Width, Length, StationType)

The 3 splits train/, valid/ and test/ can be obtained using the sumo-dataset-generator separately for each split, and then manually extracting the vmap.parquet in the shared parent folder once (since always the same).

You can refer to this schema to better understand th full work pipeline:

flowchart LR
    subgraph OFF["offline"]
      A[[XML sumo files]] --> B[["sumo-dataset-generator"]] --> C[("SUMO dataset")] -.-> D[["build.py"]]
      E[["main.py"]]
      F[["test.py"]] --> G([Virtual Metrics])
      H[["compare_predictions_gt.py"]] --> I([Online Metrics])
    end
    subgraph ON["online"]
      J[[ActiveMQ + TRACEN-X]] -.-> K[[S-LDM]]
      L[["rcv.py"]] --> M[(scores CSV)]
    end
    A -.-> J
    D -."graphs/*.pt,\n.map/vmap.pth".-> E
    E -."best_state.pth".-> F
    E -."best_state.pth".-> K
    K -."best_state.pth".-> L
    C -."labels.parquet".-> H
    C -.-> F
    M -.-> H
Loading

CLI reference

All entrypoints are based on the click python library. Run uv run <script> --help for the live option list.

build.py

Turns the parquet dataset into torch graphs. In particular:

  • for each split, packs.parquet, vinfo.parquet, and labels.parquet are used to build the graphs related to traffic: the output is a .graphs/pack_<PackId>.pt for each pack, and a single .graphs/metadata.json
  • vmap.parquet is used to build the graph associated with the map: the output is shared to all splits in .map/vmap.pth

Usage

uv run build.py <DATA_PATH> -l <ACTIVE_LABEL> [options]

Options

Option Default Required Description
-r, --radius-threshold 30.0 Max vehicle distance (m) for edge creation.
-l, --active-label - Index of the single label to consider (see LabelsEnum).
-f, --frames-num 100 Number of frames per pack.
--map-only off Only build .map/vmap.pth, skip graphs.
--map.lat-conn.max-angle 30.0 Max angle (deg) for lateral lane connections.
--map.lat-conn.proximity-threshold 1.0 Proximity (m) for lateral lane connections.
-T, --threads 1 Parallel processes for graph building.

Example

uv run build.py /data/sldm -l 1 -f 100 -T 8

Builds map + graphs for label 1 (OVERTAKE) on train/, valid/ and test/ (if present).

main.py

Training + validation. Loads precomputed graphs from <inputdir>/{train,valid}/.graphs, optional map from <inputdir>/.map/vmap.pth, z-score normalises using train statistics, then sweeps the combinations of GRUSAGE_PARAMS_DICT defined at the top of main.py. For each combination it trains GruSage, keeps the best validation snapshot, and writes:

  • <outdir>/configNN/GRUSAGE_best_state.pth (or GRUSAGE_MAP_best_state.pth with map) — snapshot embedding model ip_dict, state_dict, z-score norm_stat_dict and train_prior;
  • <outdir>/configNN/GRUSAGE_*_trev_plot.png — train/val accuracy (+ ROC AUC / precision / recall) plot.

Config dirs auto-number from the highest existing index in outdir, so multiple sweeps append to it.

Usage

uv run main.py <INPUTDIR> <OUTDIR> -l <LABEL_NUM> [options]

Options

Option Default Required Description
-l, --label-num - Label to train on (matches -l passed to build.py).
--cut None If set, cuts frames after the given index, allowing prediction at earlier timesteps.
--include-map off Feed map encoder features (requires .map/vmap.pth).
-T, --threads 1 Parallel param-sweep workers sharing the GPU (memory-bounded; keep small to avoid OOM).

Sweep hyperparameters are not CLI args: edit GRUSAGE_PARAMS_DICT at the top of main.py. Dependent values use (lambda x: ..., "other_key") tuples.

Example

uv run main.py /data/sldm ./runs -l 1 --include-map -T 2

test.py

Evaluates the best validation snapshot on the test split. The test split comes from the sumo dataset generator with a sliding window over a continuous simulation, so PackIds are sorted numerically before computing metrics to follow real simulation time. Moreover, they are overlapping, and there is no control over isolation of critical events. For this reason, it is suggested to use the test script wth the -e/--event-metrics option, that provides a measure that is more informative taking in account such difference with respect to the carefully-crafted training and validation packs.

Indeed, for each active label it writes PackMetrics (per-pack confusion matrix, precision/recall/F1, ROC AUC, PR-AUC) and, with -e/--event-metrics, EventMetrics (event-level FAR/h, event precision/recall + test_temporal_plot_lbN.png).

Usage

uv run test.py <INPUTDIR> <OUTDIR> -w <WEIGHTS_PATH> [options]

<INPUTDIR> is the test split dir containing a .graphs/ built with build.py -l <same-label>.

Options

Option Default Required Description
-w, --weights - GRUSAGE_*_best_state.pth snapshot from main.py.
-b, --batch-size 64 Inference batch size.
--threshold 0.5 Decision threshold on the sigmoid score.
--cut None Frame cut used at training (must match main.py --cut).
-e, --event-metrics off Also compute event-level metrics (FAR/h, etc.).
--sim-duration 60 Total simulation duration (s).
--calibrate-priors off Apply Bayes prior-shift calibration (deployment prior ≠ train prior).
--train-prior snap. Override train_prior for prior-shift (default: snapshot value).
--test-prior from GT Deployment P(y=1) (default: estimated from GT).
--gap-pred 5 Clustering gap (samples) for predicted events.
--gap-gt 20 Clustering gap (samples) for GT events.
--match-tol 10 Tolerance (samples) matching predicted clusters to GT events.

Example

uv run test.py /data/sldm/test ./runs/config01/t030/ --threshold 0.30 -w ./runs/config01/GRUSAGE_best_state.pth -e

rcv.py

Online inference entrypoint spawned by the S-LDM binary. It reads JSON-encoded frames from a FIFO (named pipe), uses them to create fixed-size packs, builds a graph on the fly with GraphOnlineCreator (resampling positions to vehicle center, padding missing frames, z-score normalising with the snapshot's statistics), runs GruSage on GPU, and appends the sigmoid score (or . when the pack has no vehicles) to a CSV.

Two threads cooperate in a Producer-Consumer scheme: pipeout_producer parses JSON lines off the FIFO and buffers packs into a shared deque; infer_consumer pulls a full pack, runs inference, writes one score row, then slides the window by popping the oldest pack.

rcv.py forces CUDA inference (model.cuda()); on a CPU-only host change the call.

Usage

uv run rcv.py -f <FIFO_PATH> -p <PACK_SIZE> -s <SNAPSHOT_PATH> [options]

Options

Option Default Required Description
-f, --fifo-path - Path to the FIFO (named pipe) to read from. The S-LDM creates it and writes data into it; its full name is logged on run and also directly forwarded to rcv.py when spawning it.
-p, --pack-size - Frames per pack (must match frames_num used at training).
-s, --snapshot-path - Path to .pth best snapshot from main.py.
-O, --output-csv-file out.csv CSV where streamed sigmoid scores are appended.

The script is designed to be run by the S-LDM, manual running it is discouraged!

compare_predictions_gt.py

Offline evaluation of the scores produced by rcv.py. After running the S-LDM with the inference module enabled, the CSV output is read (one row per pack, in arrival order) and aligned against a GT parquet labels.parquet sorted by numeric PackId: GT is the bitmask in MLBEncoded (decomposed with --active-label), predictions are the Scores/Score column. It then computes the same per-pack and (optional) event-level metrics as test.py.

Missing/invalid scores are handled per --nan-policy, but currently only the zero option is implemented.

Usage

uv run compare_predictions_gt.py --gt-parquet <GT.parquet> --pred-csv <PRED.csv> [options]

Options

Option Default Required Description
--gt-parquet - GT parquet with ["PackId","MLBEncoded"].
--pred-csv - Prediction CSV from rcv.py (Scores/Score column).
--threshold 0.5 Decision threshold.
--outdir cwd Where to write metrics + plots.
-e, --event-metrics off Compute event-level metrics.
--sim-duration 60 Total simulation duration (s).
--active-label None Label index to extract from MLBEncoded bitmask (e.g. 2 for TURN).
--calibrate-priors off Apply Bayes prior-shift calibration.
--train-prior - Train P(y=1) for prior-shift (else computed from --train-metadata).
--train-metadata - Train metadata.json to compute train_prior when --train-prior absent.
--test-prior from GT Deployment P(y=1) for prior-shift.
--nan-policy zero zero (set to 0.0) or drop for invalid/missing scores (unimplemented).
--gap-pred 5 Clustering gap (samples) for predicted events.
--gap-gt 20 Clustering gap (samples) for GT events.
--match-tol 10 Match tolerance (samples).

Example

uv run compare_predictions_gt.py \
    --gt-parquet /data/sldm/test/labels.parquet \
    --pred-csv ./scores.csv \
    --active-label 1 -e --outdir ./results/rcv

⚠️ When using this script to evaluate the results of a test run, the proper alignment should be guaranty, so please consider that manually running tracen-x to replay pcap traces towards the ActiveMQ broker associated with the S-LDM instance is possible, but it may be necessary to manually discard the first lines of the csv output corresponding to the delay between the startup of the S-LDM and the tracen-x script activation.


License

Distributed under the GNU General Public License v2.0 — see LICENSE.

Links

About

GNN for AI-based S-LDM trigger module

Resources

Stars

3 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages