NOTE: This is a fork of the USNet repository. The upstream repository can be found here.

The implementation is extended as follows:

• While the original has parameters, they are discarded if given via command line. I changed that and parameters given via command line are now honored.
• The test schript test.py was not fully implemented for the Cityscape dataset. I extended its functionality to support testing the model against the Cityscape dataset.
• The train script train.py saved the model and training logs into different folders. Same for the inference output of the test script test.py. The outputs are now by default all saved into the same folder. This folders name is as follows: {backbone_name}___{learning_rate}___{batch_size}___{num_epochs}___{date}.

USNet

Fast Road Segmentation via Uncertainty-aware Symmetric Network

This repository provides the implementation of USNet [arxiv] in PyTorch.

Road segmentation is significant in self-driving and mobile robot applications. USNet is proposed to achieve a trade-off between speed and accuracy in this task.

Here shows the segmentation result and the uncertainty map:

Data Preparation

KITTI Road Dataset

You may download the original images and annotations of on KITTI Road dataset from KITTI and the depth images can be found in SNE-RoadSeg. Then please setup dataset according to the following directory structure:

USNet
 |-- data
 |  |-- KITTI
 |  |  |-- training
 |  |  |  |-- calib
 |  |  |  |-- depth_u16
 |  |  |  |-- gt_image_2
 |  |  |  |-- image_2
 |  |  |-- validating
 |  |  |  |-- calib
 |  |  |  |-- depth_u16
 |  |  |  |-- gt_image_2
 |  |  |  |-- image_2
 |  |  |-- testing
 |  |  |  |-- calib
 |  |  |  |-- depth_u16
 |  |  |  |-- image_2
 |-- models
 ...

Installation

The code is developed using Python 3.7 with PyTorch 1.6.0. The code is tested using one NVIDIA 1080Ti GPU card. You can create a conda environment and install the required packages by running:

$ conda create -n usnet python=3.7
$ pip install -r requirements.txt

Training

For training USNet on KITTI Road dataset, you can run:

$ cd $USNET_ROOT
$ python train.py

When training completed, the checkpoint will be saved to ./log/KITTI_model.

Testing

Note that before testing you need to config the necessary paths or variables. Please ensure that the checkpoint exists in checkpoint_path.

To run the test on KITTI Road dataset:

$ python test.py

You can download our trained model from Google Drive or Baidu Netdisk (Code: 9zgf). The BEV-results obtained from this released model can be found in Google Drive or Baidu Netdisk (Code: csar).

If you submit this result to the KITTI benchmark, you will get a MaxF score of 96.87 for URBAN, which is similar to the reported ones in our paper.

Citation

If you find USNet useful in your research, please consider citing:

@inproceedings{Chang22Fast,
  title = {Fast Road Segmentation via Uncertainty-aware Symmetric Network},
  author = {Chang, Yicong and Xue, Feng and Sheng, Fei and Liang, Wenteng and Ming, Anlong},
  booktitle = {IEEE International Conference on Robotics and Automation (ICRA)},
  year = {2022}
}

Acknowledgement

The source code of surface normal estimator in our method follows SNE-RoadSeg, we do appreciate this great work. Besides, the code of acquiring uncertainty in our method is adapted from TMC.