We put forward an efficient multi-task network that can jointly handle three crucial tasks in autonomous driving: object detection, drivable area segmentation and lane detection to save computational costs, reduce inference time as well as improve the performance of each task. Our work is the first to reach real-time on embedded devices while maintaining state-of-the-art level performance on the BDD100K dataset.
We design the ablative experiments to verify the effectiveness of our multi-tasking scheme. It is proved that the three tasks can be learned jointly without tedious alternating optimization.
In table 4, E, D, S and W refer to Encoder, Detect head, two Segment heads and whole network. So the Algorithm (First, we only train Encoder and Detect head. Then we freeze the Encoder and Detect head as well as train two Segmentation heads. Finally, the entire network is trained jointly for all three tasks.) can be marked as ED-S-W, and the same for others.
Visualization
Traffic Object Detection Result
detect result
Drivable Area Segmentation Result
Lane Detection Result
Notes:
The visualization of lane detection result has been post processed by quadratic fitting.
Project Structure
python
1├─inference
2│ ├─images # inference images3│ ├─output # inference result4├─lib
5│ ├─config/default # configuration of training and validation6│ ├─core
7│ │ ├─activations.py # activation function8│ │ ├─evaluate.py # calculation of metric9│ │ ├─function.py # training and validation of model10│ │ ├─general.py #calculation of metric、nms、conversion of data-format、visualization11│ │ ├─loss.py # loss function12│ │ ├─postprocess.py # postprocess(refine da-seg and ll-seg, unrelated to paper)13│ ├─dataset
14│ │ ├─AutoDriveDataset.py # Superclass dataset,general function15│ │ ├─bdd.py # Subclass dataset,specific function16│ │ ├─hust.py # Subclass dataset(Campus scene, unrelated to paper)17│ │ ├─convect.py
18│ │ ├─DemoDataset.py # demo dataset(image, video and stream)19│ ├─models
20│ │ ├─YOLOP.py # Setup and Configuration of model21│ │ ├─light.py # Model lightweight(unrelated to paper, zwt)22│ │ ├─commom.py # calculation module23│ ├─utils
24│ │ ├─augmentations.py # data augumentation25│ │ ├─autoanchor.py # auto anchor(k-means)26│ │ ├─split_dataset.py # (Campus scene, unrelated to paper)27│ │ ├─utils.py # logging、device_select、time_measure、optimizer_select、model_save&initialize 、Distributed training28│ ├─run
29│ │ ├─dataset/training time # Visualization, logging and model_save30├─tools
31│ │ ├─demo.py # demo(folder、camera)32│ │ ├─test.py
33│ │ ├─train.py
34├─toolkits
35│ │ ├─depoly # Deployment of model36├─weights # Pretraining model
Requirement
This codebase has been developed with python version 3.7, PyTorch 1.7+ and torchvision 0.8+:
Update the your dataset path in the ./lib/config/default.py.
Training
You can set the training configuration in the ./lib/config/default.py. (Including: the loading of preliminary model, loss, data augmentation, optimizer, warm-up and cosine annealing, auto-anchor, training epochs, batch_size).
If you want try alternating optimization or train model for single task, please modify the corresponding configuration in ./lib/config/default.py to True. (As following, all configurations is False, which means training multiple tasks end to end).
python
1# Alternating optimization2_C.TRAIN.SEG_ONLY =False# Only train two segmentation branchs3_C.TRAIN.DET_ONLY =False# Only train detection branch4_C.TRAIN.ENC_SEG_ONLY =False# Only train encoder and two segmentation branchs5_C.TRAIN.ENC_DET_ONLY =False# Only train encoder and detection branch67# Single task 8_C.TRAIN.DRIVABLE_ONLY =False# Only train da_segmentation task9_C.TRAIN.LANE_ONLY =False# Only train ll_segmentation task10_C.TRAIN.DET_ONLY =False# Only train detection task
Start training:
python tools/train.py
Evaluation
You can set the evaluation configuration in the ./lib/config/default.py. (Including: batch_size and threshold value for nms).
You can store the image or video in --source, and then save the reasoning result to --save-dir
python tools/demo --source inference/images
Camera
If there are any camera connected to your computer, you can set the source as the camera number(The default is 0).
python tools/demo --source 0
Deployment
Our model can reason in real-time on Jetson Tx2, with Zed Camera to capture image. We use TensorRT tool for speeding up. We provide code for deployment and reasoning of model in ./toolkits/deploy.
Citation
If you find our paper and code useful for your research, please consider giving a star and citation:
BibTeX
1@misc{2108.11250,
2Author = {Dong Wu and Manwen Liao and Weitian Zhang and Xinggang Wang},
3Title = {YOLOP: You Only Look Once for Panoptic Driving Perception},
4Year = {2021},
5Eprint = {arXiv:2108.11250},
6}