LSTR: Lane Shape Prediction with Transformers

😎End-to-end architecture: Directly output lane shape parameters.
⚡Super lightweight: The number of model parameters is only 765,787.
⚡Super low complexity: The number of MACs (1 MAC = 2 FLOP) is only 574.280M.
😎Training friendly: Lower GPU memory cost. Input (360, 640, 3) with batch_size 16 uses 1245MiB GPU usages.

PyTorch(1.5.0) training, evaluating and pretrained models for LSTR (Lane Shape Prediction with Transformers). We streamline the lane detection to a single-stage framework by proposing a novel lane shape model that achieves 96.18 TuSimple accuracy.

For details see End-to-end Lane Shape Prediction with Transformers by Ruijin Liu, Zejian Yuan, Tie Liu, Zhiliang Xiong.

Updates!!

【2021/12/03】:fire: Our new work Learning to Predict 3D Lane Shape and Camera Pose from a Single Image via Geometry Constraints by Ruijin Liu, Dapeng Chen, Tie Liu, Zhiliang Xiong, Zejian Yuan is accepted by AAAI2022! The preprint paper and codes will be released soon!
【2021/11/23】 We now support Train and Test Custom Data. Tutorial: Train and Test Your Custom Data.
【2021/11/16】 We fix the Multi-GPU Training.
【2020/12/06】 We now support CULane Dataset.

Comming soon

LSTR-nano(New backbone): 96.33 TuSimple accuracy with only 40% MACs (229.419M) and 40% #Params (302,546) of LSTR.
Mosaic Augmentation.
Loguru based logger module.
Geometry based loss functions.
Segmentation prior.

Model Zoo

We provide the baseline LSTR model file (trained on TuSimple train and val sets after 500000 iterations) in the ./cache/nnet/LSTR/LSTR_500000.pkl (~3.1MB).

Data Preparation

Download and extract TuSimple train, val and test with annotations from TuSimple. We expect the directory structure to be the following:

TuSimple/
    LaneDetection/
        clips/
        label_data_0313.json
        label_data_0531.json
        label_data_0601.json
        test_label.json
    LSTR/

Set Envirionment

Linux ubuntu 16.04

conda env create --name lstr --file environment.txt

After you create the environment, activate it

conda activate lstr

Then

pip install -r requirements.txt

Training and Evaluation

To train a model:

(if you only want to use the train set, please see ./config/LSTR.json and set "train_split": "train")

python train.py LSTR

Visualized images are in ./results during training.
Saved model files (every 5000 iterations) are in ./cache during training.

To train a model from a snapshot model file:

python train.py LSTR --iter 500000

To evaluate (GPU 603MiB usage when evaluating single image iteratively), then you will see the paper's result:

python test.py LSTR --testiter 500000 --modality eval --split testing

To evaluate FPS (set --batch to maximum the FPS, GPU 877MiB usage if you repeat each image 16 times):

python test.py LSTR --testiter 500000 --modality eval --split testing --batch 16

To evaluate and save detected images in ./results/LSTR/500000/testing/lane_debug:

python test.py LSTR --testiter 500000 --modality eval --split testing --debug

Demo (displayed parameters are rounded to three significant figures.)

To evaluate and save decoder attention maps (store --debugEnc to visualize encoder attention maps):

python test.py LSTR --testiter 500000 --modality eval --split testing --debug --debugDec

To evaluate on a set of images (store your images in ./images, then the detected results will be saved in ./detections):

python test.py LSTR --testiter 500000 --modality images --image_root ./ --debug

Citation

@InProceedings{LSTR,
author = {Ruijin Liu and Zejian Yuan and Tie Liu and Zhiliang Xiong},
title = {End-to-end Lane Shape Prediction with Transformers},
booktitle = {WACV},
year = {2021}
}

License

LSTR is released under BSD 3-Clause License. Please see LICENSE file for more information.

Contributing

We actively welcome your pull requests!

Acknowledgements

DETR

PolyLaneNet

CornerNet

liuruijin17 / LSTR