HMAP Traffic Light Classifier

traffic-light status classification Stub for the HMAP / DriveWorks perception stack. Stage 1 finds traffic-light boxes (YOLO / COCO traffic light); a compact ResNet-18 + squeeze-and-excitation head classifies each letterboxed crop into Red / Yellow / Green

Built for internal HMAP integration (camera server → fusion / controller). Dataset roots (S2TLD, LISA, etc.) are not shipped in this repository; paths are configured per machine.

---

Architecture

Camera frame (RGBA, rectified)
        │
        ▼
┌───────────────────┐
│  Stage 1 — YOLO   │  COCO class 9, NMS, 640×640 (DriveWorks: yolov3_640x640.bin)
│  (detector)       │
└─────────┬─────────┘
          │ 2D boxes (traffic light)
          ▼
┌───────────────────┐     letterbox 32×96, ImageNet norm, 10% box padding
│  Stage 2b — TL    │  ResNet-18 (ImageNet) → SE block → 1×1 conv logits
│  classifier       │
└─────────┬─────────┘
          │
          ▼
   Per-box: { Red | Yellow | Green } + confidence

Stage 1 (this repo, Python / C++): YOLOv3-via-OpenCV in test_inference.py for parity with drivesegserver TensorRT; production uses the same box semantics on the Orin.

Stage 2b (this repo): Input tensor layout (1, 3, 96, 32) (batch 1, CHW). DLA-friendly ops (standard Conv2d, BN, ReLU, no nn.Linear in the logit head).

---

Results

Item	Value (typical / indicative)
TL classes	3 (Red, Yellow, Green)
Crop geometry	Letterbox 32×96 (portrait), BGR → RGB norm
Classifier params	~11.4M total (ResNet-18 backbone + SE + head)
ONNX export	Fixed batch = 1 — crop (1,3,96,32), logits (1,3)
Training	Dataset-dependent (S2TLD / LISA / folder); see outputs/*/checkpoints/best.pth
Inference (Python)	YOLOv3 + TL on GPU; ms/frame depends on image count and GPU

---

Project structure

HMAP_trafficlight_stub/
├── models/
│   └── traffic_light_classifier.py   # ResNet-18 + SE + letterbox preprocess + ONNX-oriented head
├── datasets/
│   └── traffic_light_dataset.py      # folder | LISA | S2TLD (VOC XML), aug, class weights
├── tl_driveseg_server/
│   └── drivesegserver.cu             # 4-cam server: seg + det + Stage2a + TL Stage2b + IPC
├── tl_driveseg_client/
│   └── main.cpp                      # Socket client consuming server metadata
├── train_traffic_light_classifier.py # Training loop, checkpoints, TensorBoard, two-phase LR
├── export_onnx.py                    # ONNX → DriveWorks tensorRT_optimization (batch 1)
├── two_stage_detector_with_tl.py     # Optional glue into HMAP3D-Net (needs repo on PYTHONPATH)
├── test_inference.py                 # YOLOv3 (OpenCV DNN) + TL visual sanity checks
├── harvest_tl_crops.py               # Mine TL crops from video / folders via YOLO
├── yolov3.cfg / yolov3.weights       # Darknet weights for local OpenCV test (download separately)
├── pointpillars_minimal_epoch20_static.onnx   # Optional LiDAR Stage2a (plugins + TRT on Orin)
└── outputs/                          # e.g. tl_classifier_*/checkpoints/, inference_samples/

---

Input format (Stage 2b)

Stage	Description
From YOLO	Axis-aligned box in full-image pixel coordinates (after rectification / same space as training crops).
Padding	10% margin on width/height (matches two_stage_detector_with_tl.py and tl_driveseg_server TL kernel).
Letterbox	Scale crop to fit 32×96 without stretch; pad with ImageNet mean colour.
Tensor	float32 CHW, RGB, normalized with ImageNet mean/std (same as training).

---

Training

# Setup
cd ~/HMAP_trafficlight_stub   # or your clone path
python3 -m venv .venv_tl && source .venv_tl/bin/activate
# Environment can be pinned with your own requirements.txt (torch+cu*, opencv, ultralytics if using two_stage, etc.)
pip install torch torchvision tqdm matplotlib tensorboard opencv-python-headless

# Example — S2TLD VOC layout (720p subset shown; adjust path)
python train_traffic_light_classifier.py \
    --data-root "S2TLD（720x1280）" \
    --fmt s2tld \
    --epochs 30 \
    --batch-size 64 \
    --output-dir outputs/tl_classifier_run1

---

ONNX & DriveWorks (Orin)

Export (fixed batch 1 — required for tensorRT_optimization without dynamic profiles):

python export_onnx.py \
    --checkpoint outputs/tl_classifier_run1/checkpoints/best.pth \
    --output tl_classifier.onnx \
    --verify

On the target (DriveWorks SDK tools/dnn):

./tensorRT_optimization --modelType=onnx \
    --onnxFile=/path/to/tl_classifier.onnx \
    --out=tl_classifier.bin \
    --half2=1

Load the resulting .bin with dwDNN (see tl_driveseg_server --tl_model).

---

Python test (YOLOv3 + TL)

source .venv_tl/bin/activate
python test_inference.py \
    --images path/to/frames \
    --tl-ckpt outputs/tl_classifier_run1/checkpoints/best.pth \
    --yolo-cfg yolov3.cfg \
    --yolo-weights yolov3.weights \
    --output outputs/inference_samples

---

Deployment target

NVIDIA Drive AGX Orin — Rough dataflow:

Cameras → DriveSeg / perception → YOLO boxes → [TL classifier .bin] → downstream (fusion, HUD, logging)

References

• S2TLD — SJTU small traffic light dataset: Thinklab-SJTU/S2TLD
• LISA — Traffic-light imagery / annotations (e.g. GTDLBench — LISA traffic dataset)
• Two-stage TL (bulb-level) — Detection and recognition framework: ResearchGate publication
• NVIDIA DriveWorks — DriveWorks documentation (DNN, TensorRT optimization on target)

---