# 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](https://git.hipert.unimore.it/mverucchi/shelfnet) 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

[![tkDNN ShelfNet18 on MASA](https://img.youtube.com/vi/oiMzOLjo47A/0.jpg)](https://www.youtube.com/watch?v=oiMzOLjo47A "tkDNN ShelfNet18 on MASA")

## Existing tests and supported networks

| Test Name         | Network                                       | Dataset                                                       | N Classes | Input size    | Weights                                                                   |
| :---------------- | :-------------------------------------------- | :-----------------------------------------------------------: | :-------: | :-----------: | :------------------------------------------------------------------------ |
| shelfnet              | ShelfNet18_realtime<sup>1</sup>                           | [Cityscapes](https://www.cityscapes-dataset.com/)                          | 19        | 1024x1024       | [weights](https://cloud.hipert.unimore.it/s/mEDZMRJaGCFWSJF/download)                                                                   |
| shelfnet_berkeley              | ShelfNet18_realtime<sup>1</sup>                           | [DeepDrive](https://bdd-data.berkeley.edu/)                          | 20        | 1024x1024       | [weights](https://cloud.hipert.unimore.it/s/m92e7QdD9gYMF7f/download)                                                                   |
| shelfnet_mapillary              | ShelfNet18_realtime<sup>1</sup>                           | [Mapillary Vistas](https://www.mapillary.com/dataset/vistas?pKey=aFWuj_m4nGoq3-tDz5KAqQ)<sup>*</sup>                          | 15        | 1024x1024       | [weights](https://cloud.hipert.unimore.it/s/6WnZCKLjik7xrny/download)                                                                   |

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.