tkDNN/README_seg.md

Semantic Segmentation with tkDNN

Currently tkDNN supports only ShelfNet as semantic segmentation network.

Export weights from Shelfnet

To get the weights needed to run Shelfnet tests use this fork of a Pytorch implementation of Shelfnet network.

git clone https://git.hipert.unimore.it/mverucchi/shelfnet
cd shelfnet 
cd ShelfNet18_realtime
conda env create --file shelfnet_env.yml
conda activate shelfnet
mkdir layer debug
python export.py

Run the demo

To run the semantic segmentation demo follow these steps (example with shelfnet_mapillary):

rm shelfnet_mapillary_fp32.rt        # be sure to delete(or move) old tensorRT files
export TKDNN_BATCHSIZE=4             # be sure you have batch size > than 1 if you want to run inference on images bigger than 1024
./test_shelfnet_mapillary            # run the yolo test (is slow)
./seg_demo shelfnet_mapillary_fp32.rt ../demo/yolo_test.mp4 1 15

In general the demo program takes the following parameters:

./seg_demo <network-rt-file> <path-to-video> <n-batches> <number-of-classes> <resize-flag> <baseline-resize> <show-flag> <write-pred>

where

• <network-rt-file> is the rt file generated by a test
• <<path-to-video> is the path to a video file or a camera input
• <n-batches> number of batches to use in inference (N.B. you should first export TKDNN_BATCHSIZE to the required n_batches and create again the rt file for the network).
• <number-of-classes>is the number of classes the network is trained on
• <resize-flag> if set to 0 the demo will not resize the input frames, but use it as it is, otherwise it will resize it.
• <baseline-resize> is <resize-flag> is set to 1, then the input frames will be proportionally resized using <baseline-resize> as width baseline.
• <show-flag> if set to 0 the demo will not show the visualization but save the video into result.mp4 (if n-batches ==1)
• <write-pred> if set to 0 (deafult) the demo will run, otherwise the evaluation of a dataset will run and the output of the segmentation will be saved. Attention: this is under development and paths are embedded, so change them in the code in advance.

N.b. By default it is used FP32 inference

Existing tests and supported networks

Test Name	Network	Dataset	N Classes	Input size	Weights
shelfnet	ShelfNet18_realtime1	Cityscapes	19	1024x1024	weights
shelfnet_berkeley	ShelfNet18_realtime1	DeepDrive	20	1024x1024	weights
shelfnet_mapillary	ShelfNet18_realtime1	Mapillary Vistas*	15	1024x1024	weights

1. Zhuang, Juntang, et al. "ShelfNet for fast semantic segmentation." Proceedings of the IEEE International Conference on Computer Vision Workshops. 2019.

*. Mapillary Vistas has originally 66 classes, but we reduced them to 15 to improve the results on the categories of our interest.