Traffic Signal Optimization

Turns raw intersection camera frames into vehicle density maps, forecasts those maps 15, 30 and 60 minutes ahead, and trains a multi-agent RL controller in SUMO to pick green phases across a 17-signal city grid.

Three stages, each a separate PyTorch model with its own train and test script:

Stage	Model	Input	Output
1. Density estimation	VehicleEstimationModel (grouped-conv residual CNN, about 0.2M params)	72x72 RGB frame	18x18 density map, map.sum() = vehicle count
2. Flow prediction	ST3DNet x3 (3D-conv closeness + weekly streams)	6 recent maps + 4 same-time-previous-weeks maps	density map 15 min / 30 min / 1 hr ahead
3. Signal control	HierarchicalSignalNet trained with MAPPO	per-approach 18x18 occupancy and queue grids from SUMO	next green phase per junction

Status: research code. Stage 1 is trained on real Jaipur CCTV frames and its checkpoint is in the repo. Stage 2 is trained on synthetic data. Stage 3 is trained and evaluated in simulation only. See Limitations before relying on any stage.

Why

Loop detectors and point sensors give one number per lane. A camera frame shows the whole approach, but object detection loses count when vehicles overlap, which is the normal case in the mixed two-wheeler traffic in the Jaipur frames here. This repo regresses a density map instead: each of the 18x18 cells holds a fractional vehicle count, the sum is the total, and the spatial layout is kept.

That one representation is the data format for all three stages. The forecaster (Stage 2) reads stacks of density maps, and the RL controller (Stage 3) reads 18x18 grids per approach, so there is one feature format rather than three.

Stage 3 learns coordination between junctions with a hierarchical network (approach to junction to network) and one centralised critic, rather than an independent agent per light. Junction adjacency is read from the SUMO network file at start-up, so the model shape follows whatever .net.xml you give it.

Demo

Stage 1 training data: one of the 40 frames from the Jaipur (Rambagh) CCTV feed. The matching file data/annotation/00h_00m_00s.json marks 26 vehicle centre points, and preprocess.py turns them into the 18x18 ground-truth map.

Quickstart

All scripts use paths relative to the repo root (data/..., checkpoints/...), so run everything from the root.

1. Install (Stages 1 and 2)

git clone <repo-url> traffic-signal-optimizer && cd traffic-signal-optimizer
python -m venv env && source env/bin/activate
pip install -r requirements.txt       # pinned: torch 2.10, numpy, scipy, pillow, matplotlib

2. Run the shipped density estimator

checkpoints/estimation/best_model.pt is tracked in git, so this works with no training:

python test/test_estimation.py
# prints the map range and vehicle count for data/images/0000.jpg
# saves data/output/estimation/jaipur_<timestamp>.png

3. Train the forecasters, then run the camera pipeline

Stage 2 checkpoints are not in the repo. train_prediction.py fits one ST3DNet per horizon on synthetic data (200 samples, 20 epochs each, CPU is fine):

python train/train_prediction.py      # writes checkpoints/prediction/st3dnet_{15min,30min,1hr}.pt
python test/test_pipeline.py          # seeds 4 weekly + 6 recent maps, then estimates and forecasts for 2026-02-11 08:00

Run step 3 before test_pipeline.py: inference.Pipeline skips a missing checkpoint silently and would forecast with random weights.

4. Signal control in SUMO (Stage 3)

pip install eclipse-sumo sumo-rl      # SUMO binaries, sumolib, traci, plus the gymnasium env wrapper
python sumo_config/generate_network.py   # optional: the generated city_*.xml files are already committed
python train/train_signal.py          # MAPPO, 200 updates x 128 steps on CPU; writes checkpoints/signal/mappo_signal.pt
python test/test_signal.py            # 3 episodes with the trained agent, then 3 with random actions
python test/visualize_simulation.py   # 600 s fixed-time SUMO run -> data/output/simulation.mp4 (needs ffmpeg)

train_signal.py, test_signal.py and visualize_simulation.py set SUMO_HOME themselves from the eclipse-sumo package. Keep the venv activated: generate_network.py shells out to netconvert and visualize_simulation.py launches sumo, and both executables come from the eclipse-sumo wheel. The equivalent without activating the venv:

export SUMO_HOME=$(env/bin/python -c "import sumo; print(sumo.SUMO_HOME)")
env/bin/python train/train_signal.py

5. Retrain Stage 1 (optional)

python preprocess.py                  # VIA point annotations -> data/density_maps/00h_*.npy (already committed)
python train/train_estimation.py      # 5-fold CV on the 40 Jaipur frames; picks mps, then cuda, then cpu

No Dockerfile or compose file exists.

How it works

 CAMERA PATH (inference.Pipeline)                        SIMULATION PATH (SUMO)
 ================================                        ======================

 data/image/*.jpg      data/annotation/*.json            sumo_config/generate_network.py
 720x480 CCTV frame    VIA vehicle centre points         32 nodes, 17 signals, 16 flows
        |                       |                                    | netconvert
        |                  preprocess.py                             v
        |                  Gaussian splat, sum = count    city_network.net.xml + city_routes.rou.xml
        |                       |                                    |
        |                       v                                    v
        |              data/density_maps/*.npy            NetworkEnv (architecture/signal_optimizer.py)
        |                       |                         gymnasium wrapper over sumo-rl
        v                       v                         lane occupancy + halting -> per approach
 +---------------------------------------+                (6, 18, 18, 2) grid, phase one-hot,
 | Stage 1  VehicleEstimationModel       |                min-green flag
 | architecture/vehicle_estimation_model |                           |
 | 72x72 RGB -> 18x18 map, count = sum   |                           v
 +------------------+--------------------+                +-------------------------------------+
                    | every frame                         | Stage 3  HierarchicalSignalNet      |
                    v                                     |  L1 patch CNN + GCN per approach    |
 data/density_history/<location>/<ts>.npy                 |  L2 set-attention pooling per node  |
   closeness: 6 maps, 15 min apart                        |  L3 graph transformer + GRU         |
   weekly:    4 maps, 1 week apart                        |  L4 actor (phase) + CTDE critic     |
                    |                                     |  PPO + GAE (train/train_signal.py)  |
                    v                                     +------------------+------------------+
 +---------------------------------------+                                   |
 | Stage 2  ST3DNet x3                   |                                   v
 | architecture/flow_prediction.py       |                green phase per junction -> SUMO
 | 15 min / 30 min / 1 hr density maps   |                (+ emergency override at p95 occupancy)
 +---------------------------------------+

The two paths are not joined yet: Stage 3 reads lane statistics from SUMO, not the Stage 1 maps or Stage 2 forecasts (see Limitations).

1. preprocess.py reads each VIA JSON, scales every vehicle point from 720x480 to the 18x18 grid, splats a Gaussian whose sigma is half the mean distance to the 3 nearest neighbours (clamped to 0.8..2.5 cells), and rescales the map so it sums to the annotated count.
2. train/train_estimation.py trains VehicleEstimationModel (two 7x7 conv stem, two grouped-conv residual blocks, 1x1 regression head with ReLU output) on 72x72 resized frames. Loss is utils.DensityMapLoss: pixel MSE plus 2.0 x squared count error. 5 folds, 10x augmentation, early stopping at 120 epochs without improvement. The five fold weights are averaged into best_model.pt.
3. inference.Pipeline.run(image, timestamp) estimates a map, saves it as data/density_history/<location>/<YYYYmmdd_HHMMSS>.npy, then looks for 6 maps at 15-minute steps (10-minute tolerance) and 4 maps at 1-week steps (30-minute tolerance). If both sets exist it runs the three ST3DNets; otherwise it returns a missing list instead of a forecast.
4. ST3DNet runs a 3x3x3 Conv3d and two residual units over the closeness stack and a (4,1,1) Conv3d over the weekly stack, weights each with a learned per-pixel RecalibrationBlock, and fuses them with two learned 18x18 matrices under tanh.
5. sumo_config/generate_network.py writes 32 nodes (20 junctions, 17 signalised, 12 entry and exit points), three road types (arterial 3 lanes 50 km/h, collector 2 lanes 40 km/h, local 1 lane 30 km/h), 82 directed edges and 16 hourly flows totalling 5010 veh/h, then calls netconvert.
6. NetworkEnv wraps sumo_rl.SumoEnvironment (10 s decision step, 45 s min green, 120 s max green, diff-waiting-time reward). discover_topology() reads the net at start-up: signals with 2 or more green phases are controllable, lanes are grouped by incoming edge into approaches, and signals that share a road become graph neighbours.
7. HierarchicalSignalNet turns each approach's last 6 grids into 10 tokens (9 spatial regions + 1 interaction token), lets sibling approaches exchange interaction tokens through cross-attention, pools each junction to 4 tokens with a Set Transformer (SAB + PMA), runs a graph transformer with a GRUCell and 8-frame temporal cross-attention over junctions, then HybridActor samples a phase (Categorical) and a duration (Normal, clamped to 47..120 s) and CTDECritic scores the global state.
8. train/train_signal.py collects 128-step rollouts, computes GAE, runs 4 clipped-PPO epochs per update for 200 updates, and anneals the queue-variance penalty (beta, 0 to 0.2) and the neighbour-spillback penalty (gamma_coord, 0 to 0.1) over training. The best-reward checkpoint goes to checkpoints/signal/mappo_signal.pt.

Results in the repo

The only logged metric in the repo is inside the tracked Stage 1 checkpoint. train_estimation.py stores each fold's best validation MAE (vehicles per frame, 8 held-out frames per fold) in best_model.pt under fold_maes:

Fold	0	1	2	3	4	Mean
Val MAE (vehicles)	3.62	2.00	2.67	3.48	4.82	3.32 +/- 0.95

For scale, the first annotated frame holds 26 vehicles. Stage 2 and Stage 3 checkpoints and evaluation logs are not committed, so there are no forecast-error or waiting-time numbers to report. Once trained, test/test_signal.py prints episode reward and system_total_waiting_time for the agent and for a random-action baseline.

Features

Feature	Where
Ground-truth density maps from VIA point annotations with adaptive Gaussian sigma and exact count preservation	preprocess.py
Count-aware loss: pixel MSE plus weighted squared count error	utils.DensityMapLoss
5-fold CV with 10x augmentation (flip, brightness, contrast, saturation, noise, blur, crop) and weight-averaged final model	train/train_estimation.py
Two-stream ST3DNet (closeness + weekly) with learned per-pixel recalibration, one model per horizon	architecture/flow_prediction.py, train/train_prediction.py
Timestamped density history store with nearest-file lookup and tolerance; forecast only when history is complete	inference.Pipeline
Procedural SUMO city: 32 nodes, 17 signals, 82 edges, 16 routed flows	sumo_config/generate_network.py
Gymnasium env over sumo-rl exposing per-approach (6, 18, 18, 2) grids, phase one-hot and min-green flag	NetworkEnv
Runtime topology discovery: controllable signals, approaches per junction, junction adjacency	discover_topology()
Four-level policy: patch GCN, sibling cross-attention, Set Transformer pooling, graph transformer + GRU, hybrid actor, CTDE critic	HierarchicalSignalNet
Reward shaping: waiting-time delta minus annealed queue-variance and neighbour-spillback penalties	NetworkEnv._compute_reward, train_signal.anneal
Emergency override: any controllable signal with a lane at or above the 95th-percentile occupancy is forced to phase 0	NetworkEnv._check_emergency
Simulation renderer: edge colour by occupancy, vehicles by speed, signal state, stitched to MP4 with ffmpeg	test/visualize_simulation.py

Configuration

There are no CLI flags or config files. Settings are module-level constants; edit the file to change them. SUMO_HOME is the only environment variable and the Stage 3 scripts set it themselves.

Setting	File	Default	Purpose
ESTIMATION_CKPT, PREDICTION_CKPTS	inference.py	checkpoints/estimation/best_model.pt, checkpoints/prediction/st3dnet_{15min,30min,1hr}.pt	Weights loaded by Pipeline; a missing file is skipped without error
T_C, T_W	inference.py	6, 4	Recent maps (15 min apart) and weekly maps (1 week apart) needed before forecasting
HISTORY_DIR	inference.py	data/density_history	Per-location store of every estimated map
IMG_W, IMG_H, MAP_H, MAP_W	preprocess.py	720x480, 18x18	Annotation frame size and density map size
NUM_FOLDS, EPOCHS_PER_FOLD, PATIENCE	train/train_estimation.py	5, 600, 120	Cross-validation schedule
BATCH_SIZE, LR, WEIGHT_DECAY, DROPOUT, COUNT_WEIGHT	train/train_estimation.py	8, 2e-4, 1e-3, 0.25, 2.0	AdamW and loss weighting
HORIZONS	train/train_prediction.py	15min=1, 30min=2, 1hr=4 steps	One ST3DNet and checkpoint per horizon
num_samples, num_epochs, batch_size, lr	train/train_prediction.py	200, 20, 16, 1e-3	Synthetic training run
NUM_SECONDS, DELTA_TIME, MIN_GREEN, MAX_GREEN	train/train_signal.py, test/test_signal.py	3600, 10, 45, 120	Episode length, decision step and green limits in seconds
LR, GAMMA, GAE_LAMBDA, CLIP_EPS, ENTROPY_COEF, VALUE_COEF, MAX_GRAD_NORM	train/train_signal.py	3e-4, 0.99, 0.95, 0.2, 0.01, 0.5, 0.5	PPO hyperparameters
N_STEPS, N_EPOCHS, MINIBATCH_SIZE, TOTAL_UPDATES	train/train_signal.py	128, 4, 32, 200	Rollout length and update schedule
BETA_START/END, GAMMA_COORD_START/END	train/train_signal.py	0 to 0.2, 0 to 0.1	Annealed reward penalties
D_MODEL, N_HEADS, T_HIST, T_OBS, PATCH_SIZE	architecture/signal_optimizer.py	64, 4, 6, 8, 3	Token width, attention heads, grid frames per approach, temporal buffer length, patch size
EMERGENCY_DENSITY_PCTILE	architecture/signal_optimizer.py	0.95	Occupancy percentile that triggers the override
location, input_img	test/test_estimation.py, test/test_pipeline.py	jaipur, data/images/0000.jpg	History folder name and demo frame
SIM_SECONDS, FRAME_INTERVAL, FPS	test/visualize_simulation.py	600, 5, 5	Render length and frame rate

Design decisions and trade-offs

• Density regression, not detection. The count is map.sum(), so DensityMapLoss adds a squared count term on top of pixel MSE, and preprocess.py renormalises each ground-truth map to the exact annotated count. Overlapping vehicles still contribute mass.
• Small-data regime by design. 40 frames from one camera is all Stage 1 has, so training uses 5-fold CV, 10x augmentation oversampling, early stopping, and averages the five fold weights instead of picking one fold. The fold MAEs above are the honest ceiling on that data.
• 72x72 input, 18x18 output, about 0.2M parameters. The stem downsamples 4x through two pools and a stride-2 residual block. The checkpoint is 0.8 MB and the model runs on CPU, at the cost of any fine spatial detail.
• History is a folder of .npy files, not a database. Pipeline._find_nearest globs the location folder and picks the closest timestamp within tolerance. Simple, but linear in the number of stored maps.
• Stage 2 is trained on synthetic hotspots. create_dummy_prediction_data builds four Gaussian blobs with a sinusoidal trend because there is no real multi-week density history yet. The architecture is real; the weights are placeholders.
• Reward curriculum. train_signal.py starts with pure waiting-time reward and anneals in the queue-variance and neighbour-spillback penalties, so the agent learns to clear queues before it learns to balance them.
• Graph attention is fully connected. GraphTransformerLayer builds an adjacency mask from edge_index but the mask is disabled (if False). Junction adjacency is used only by the reward's spillback penalty.
• Topology is discovered at runtime and stored in the checkpoint. n_nodes, max_k, max_phases and controllable_idx are saved alongside the weights, so test_signal.py rebuilds the exact model shape for the network it was trained on.

Project layout

.
├── architecture/
│   ├── vehicle_estimation_model.py   # Stage 1 CNN
│   ├── flow_prediction.py            # Stage 2 ST3DNet
│   └── signal_optimizer.py           # Stage 3 HierarchicalSignalNet + NetworkEnv (sumo-rl wrapper)
├── train/                            # train_estimation.py, train_prediction.py, train_signal.py
├── test/                             # one runnable smoke or eval script per stage, plus visualize_simulation.py
├── inference.py                      # Pipeline: estimate -> store history -> forecast
├── preprocess.py                     # VIA annotations -> ground-truth density maps
├── utils.py                          # DensityMapLoss (used) and masked MAE/MSE helpers (unused)
├── sumo_config/                      # generate_network.py and the generated city_*.xml; older 9-signal network.net.xml is unreferenced
├── data/
│   ├── image/                        # 40 Jaipur CCTV frames, 720x480, about 30 s apart (Stage 1 training set)
│   ├── annotation/                   # VIA point annotations, one JSON per frame
│   ├── density_maps/                 # 18x18 ground-truth maps (.npy) for both image sets
│   ├── images/                       # 200 highway frames, 300x168, used only as the demo input
│   └── density_history/jaipur/       # 11 maps written by test_pipeline.py (tracked)
├── checkpoints/estimation/           # best_model.pt (weight-averaged) and fold_0..4.pt (tracked)
├── requirements.txt                  # pinned Stage 1 and 2 dependencies; SUMO stack is installed separately
└── abstract.pdf, report.pdf          # write-ups of the approach

.gitignore lists checkpoints/ and data/density_history/, but the files above were added anyway. New checkpoints you train will not be tracked unless you force-add them.

Development

There is no test framework, linter config or CI workflow. The test/ scripts are runnable checks, not a pytest suite:

python test/test_estimation.py    # Stage 1 on one frame, writes a heatmap
python test/test_prediction.py    # Stage 2 on random tensors, checks the three checkpoints load and run
python test/test_pipeline.py      # Stages 1 + 2 end to end with seeded history
python test/test_signal.py        # Stage 3 agent vs random baseline in SUMO

Outputs go to data/output/ (gitignored). Device selection: train_estimation.py prefers mps, then cuda, then cpu; train_prediction.py prefers cuda; train_signal.py and test_signal.py are pinned to cpu.

Roadmap and limitations

Evidenced by the code as it stands:

• Stages are not wired together. NetworkEnv builds its grids from SUMO lane occupancy and halting counts (_approach_to_grid), not from Stage 1 density maps, and nothing feeds Stage 2 forecasts into the agent's state.
• The duration action is never applied. HybridActor samples a green duration and PPO trains its log-prob, but NetworkEnv.step reads action['duration'] and discards it; only the phase index reaches SUMO.
• Adjacency-masked attention is built but disabled in GraphTransformerLayer.
• Stage 2 forecasts are not meaningful yet. ST3DNet is trained only on synthetic data, and its tanh output is bounded to [-1, 1] while the Stage 1 maps it receives at inference sum to the vehicle count. Training on real density history and fixing the output scale are both needed.
• Tiny, single-camera dataset. 40 annotated frames from one Jaipur junction. The 200 highway frames in data/images/ have density maps but are excluded from training by the 00h_* glob.
• Only a random-action baseline exists in test_signal.py; there is no fixed-time or Webster baseline in code.
• utils.py carries masked MAE/MSE losses and horizon metrics that no script imports.

Contributing

Open an issue or PR. Keep changes runnable from the repo root and update the constants table above if you add a setting.

License

No license file yet.