Embedded Uncertainty for Semantic Segmentation

This repo contains code for the paper: Uncertainty in Real-Time Semantic Segmentation on Embedded Systems

arXiv Link to paper.

Corrections made in arXiv version

If you came from the CVPR version of the paper, would recommend referring to the arXiv version as I found an error that I have since corrected. Specifically, I didn't square the standard deviation in Equation 15, so that what I was referring to as a variance was actually a standard deviation. Is one of those typos/mistakes that doesn't change any of the results, methods, or conclusions, but is still an error and would like to make the best efforts to correct it. I have informed the Computer Vision Foundation and Errata has been added to the CVPR repository notes here.

Description

This work expands upon the BiSeNet models for real-time semantic segmentation to include a probabilistic classifier block that allows for single pass uncertainty estimates. This has been demonstrated to permit real-time semantic segmentation on Jetson embedded GPU devices from NVIDA.

This code is based of the implementation of BiSeNet by CoinCheung, which includes instructions for compilation to using TensorRT. This code has also been modified to permit the use of the proposed probabilistic segmentation module. Instructions for TensorRT compilation can be found here.

Training

Training builds upon the pre-existing weights from here. Links to the pre-trained models are as follows,

Cityscapes

bisenetv1 download bisenetv2 download |

CocoStuffs

bisenetv1 download bisenetv2 download

ADE20k

bisenetv1 download bisenetv2 download

Running training for embedded uncertainty

Training is done using the pytorch Torchrun to allow for distributed training when available. To perform fine tuning to allow addition, you can run the tools/train_amp.py script.

An example

torchrun  --nnodes=1 --nproc_per_node=1 --rdzv_id=100 --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:29400  tools/train_amp.py --config ./configs/bayes_bisenetv2_city.py --finetune-from ./models/model_final_v2_city.pth

Where bayes_bisenetv2_city.py is a configuration dictionary with training info, and ./models/model_final_v2_city.pth is a reference to the pretrained model downloaded as described above. Descriptions for preparing datsets can be found below.

Evaluation

For evaluation, the ./tools/evaluate.py script is used. An example is provided as,

python tools/evaluate.py --weight-path ./res/bisenetv2_CocoStuff_bayes_model_final.pth --config ./configs/bayes_bisenetv2_coco.py --mode eval_bayes_prob

The --mode argument instructs how the model will operate. If --mode eval_bayes_prob, the chosen model will be evaluated using the proposed probabilistic module. If --mode eval, it will be evaluated as a normal point-estimate network and will not output uncertainty information.

Uncertainty Visualisation

For visualisation, the ./tools/demo.py script is used. An example is provided as,

python tools/demo.py --config ./configs/bisenetv2_city.py --weight-path ./res/model_final.pth  --img-path /home/ethan/exp_data/cityscapes/leftImg8bit/train/cologne/cologne_000000_000019_leftImg8bit.png 

prepare dataset

These instructions are laregely the same as from CoinCheung, but are listed here for completeness. 1.cityscapes

Register and download the dataset from the official website. Then decompress them into the datasets/cityscapes directory:

$ mv /path/to/leftImg8bit_trainvaltest.zip datasets/cityscapes
$ mv /path/to/gtFine_trainvaltest.zip datasets/cityscapes
$ cd datasets/cityscapes
$ unzip leftImg8bit_trainvaltest.zip
$ unzip gtFine_trainvaltest.zip

2.cocostuff

Download train2017.zip, val2017.zip and stuffthingmaps_trainval2017.zip split from official website. Then do as following:

$ unzip train2017.zip
$ unzip val2017.zip
$ mv train2017/ /path/to/BiSeNet/datasets/coco/images
$ mv val2017/ /path/to/BiSeNet/datasets/coco/images

$ unzip stuffthingmaps_trainval2017.zip
$ mv train2017/ /path/to/BiSeNet/datasets/coco/labels
$ mv val2017/ /path/to/BiSeNet/datasets/coco/labels

$ cd /path/to/BiSeNet
$ python tools/gen_dataset_annos.py --dataset coco

3.ade20k

Download ADEChallengeData2016.zip from this website and unzip it. Then we can move the uncompressed folders to datasets/ade20k, and generate the txt files with the script I prepared for you:

$ unzip ADEChallengeData2016.zip
$ mv ADEChallengeData2016/images /path/to/BiSeNet/datasets/ade20k/
$ mv ADEChallengeData2016/annotations /path/to/BiSeNet/datasets/ade20k/
$ python tools/gen_dataset_annos.py --ade20k

4.custom dataset

If you want to train on your own dataset, you should generate annotation files first with the format like this:

munster_000002_000019_leftImg8bit.png,munster_000002_000019_gtFine_labelIds.png
frankfurt_000001_079206_leftImg8bit.png,frankfurt_000001_079206_gtFine_labelIds.png
...

Each line is a pair of training sample and ground truth image path, which are separated by a single comma ,.

I recommand you to check the information of your dataset with the script:

$ python tools/check_dataset_info.py --im_root /path/to/your/data_root --im_anns /path/to/your/anno_file

This will print some of the information of your dataset.

Then you need to change the field of im_root and train/val_im_anns in the config file. I prepared a demo config file for you named bisenet_customer.py. You can start from this conig file.

Compilation with TensorRT

Instructions for compilation with TensorRT can be found here

Questions?

Feel free to email me (Ethan) at ej.goan@qut.edu.au

Citation

@article{goan2023uncertainty,
  title={Uncertainty in Real-Time Semantic Segmentation on Embedded Systems},
  author={Goan, Ethan and Fookes, Clinton},
  journal={arXiv preprint arXiv:2301.01201},
  year={2023}
}