Fix shelfnet berkeley, update readme, minors

Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>

demo/demo/seg_demo.cpp

@@ -39,7 +39,7 @@ int main(int argc, char *argv[]) {
     std::string net = "shelfnet_fp32.rt";
     if(argc > 1)
         net = argv[1]; 
-    std::string input = "../../ShelfNet/ShelfNet18_realtime/data/leftImg8bit/test/modena/000302.png";
+    std::string input = "../demo/yolo_test.mp4";
     if(argc > 2)
         input = argv[2]; 
     int n_batch = 1;
@@ -68,9 +68,12 @@ int main(int argc, char *argv[]) {
 
     std::string net_name;
     removePathAndExtension(net, net_name);
-    bool mapillary_15 = false; //TODO change me pls
+    bool mapillary_15_colormap = false; 
-    if(n_classes == 15 && net_name == "shelfnet_mapillary_fp32") 
+    if( n_classes == 15 && 
-        mapillary_15 = true;
+        (   net_name == "shelfnet_mapillary_fp32" || 
+            net_name == "shelfnet_mapillary_fp16" || 
+            net_name == "shelfnet_mapillary_int8"   )   ) 
+        mapillary_15_colormap = true;
 
     //net initialization 
     tk::dnn::SegmentationNN segNN;
@@ -127,7 +130,7 @@ int main(int argc, char *argv[]) {
             width = frame.cols;
 
             //inference
-            segNN.updateOriginal(frame, true, mapillary_15);
+            segNN.updateOriginal(frame, true, mapillary_15_colormap);
             if(show)
                 segNN.draw();
 

include/tkDNN/NetworkViz.h

@@ -6,7 +6,7 @@
 namespace tk { namespace dnn {
 
 cv::Mat vizFloat2colorMap(cv::Mat map, double min=0, double max=0, bool mapillary_15=false);
-cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int imgdim, double min=0, double max=0, bool mapillary_15=false);
+cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int img_h, int img_w, double min=0, double max=0, bool mapillary_15=false);
 cv::Mat vizLayer2Mat(tk::dnn::Network *net, int layer, int imgdim = 1000);
     
 }}

include/tkDNN/SegmentationNN.h

@@ -112,7 +112,7 @@ class SegmentationNN {
             cv::Mat colored;
 
             if(appy_colormap)
-                colored = vizData2Mat(tmpOutData_h, vdim, 1024, 0, classes, mapillary_15);
+                colored = vizData2Mat(tmpOutData_h, vdim, netRT->input_dim.h, netRT->input_dim.w, 0, classes, mapillary_15);
             else{
                 cv::Mat colored_fp32 (cv::Size(odim.w, odim.h),CV_32FC1, tmpOutData_h);
                 colored_fp32.convertTo(colored, CV_8UC1);
@@ -347,7 +347,7 @@ class SegmentationNN {
                     cv::Mat colored;
 
                     if(apply_colormap)
-                        colored = vizData2Mat(tmpOutData_h, vdim, 1024, 0, classes, mapillary_15);
+                        colored = vizData2Mat(tmpOutData_h, vdim, netRT->input_dim.h, netRT->input_dim.w, 0, classes, mapillary_15);
                     else{
                         cv::Mat colored_fp32 (cv::Size(odim.w, odim.h),CV_32FC1, tmpOutData_h);
                         colored_fp32.convertTo(colored, CV_8UC1);

readme/README_seg.md

@@ -18,12 +18,12 @@ python export.py
 
 ## Run the demo
 
-To run the semantic segmentation demo follow these steps (example with shelfnet_mapillary):
+To run the semantic segmentation demo follow these steps (example with shelfnet):
 ```
-rm shelfnet_mapillary_fp32.rt        # be sure to delete(or move) old tensorRT files
+rm shelfnet_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
+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)
+./test_shelfnet            # run the yolo test (is slow)
-./demo shelfnet_mapillary_fp32.rt ../demo/yolo_test.mp4 1 15
+./demo shelfnet_fp32.rt ../demo/yolo_test.mp4 1 19
 ```
 In general the demo program takes the following parameters:
 ```
@@ -37,11 +37,13 @@ where
 *  ```<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.
+*  ```<write-pred>``` if set to 0 (default) 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
 
-<!-- TODO: add gif -->
+
+![gif](output.gif "Results on yolo_test.mp4") 
+
 
 ## Existing tests and supported networks
 
@@ -49,8 +51,11 @@ N.b. By default it is used FP32 inference
 | :---------------- | :-------------------------------------------- | :-----------------------------------------------------------: | :-------: | :-----------: | :------------------------------------------------------------------------ |
 | 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.
+
+## Known issues
+
+When creating the rt file all the checks returns errors. It is due to a different resize function and handling of the original ShelfNet outputs.
+However, the network is supposed to work. 

src/NetworkViz.cpp

@@ -113,7 +113,7 @@ cv::Mat vizFloat2colorMap(cv::Mat map,double min, double max, bool mapillary_15)
     return falseColorsMap;
 }
 
-cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int imgdim, double min, double max, bool mapillary_15) {
+cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int img_h, int img_w, double min, double max, bool mapillary_15) {
     dnnType *data = nullptr;
 
     // copy to CPU
@@ -135,8 +135,7 @@ cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int imgdim, doub
         raw.copyTo(grid.rowRange(r*dim.h, r*dim.h + dim.h).colRange(c*dim.w, c*dim.w + dim.w));
     }
 
-    float ar = float(dim.w)/dim.h;
+    cv::Size vdim(img_w, img_h);
-    cv::Size vdim(ar*imgdim, imgdim);
     cv::Mat viz;
     cv::resize(grid, viz, vdim, 0, 0, 0);
     
@@ -150,7 +149,7 @@ cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int imgdim, doub
 cv::Mat vizLayer2Mat(tk::dnn::Network *net, int layer, int imgdim) {
     if(layer >= net->num_layers)
         FatalError("Could not viz layer\n");
-    return vizData2Mat(net->layers[layer]->dstData, net->layers[layer]->output_dim, imgdim);
+    return vizData2Mat(net->layers[layer]->dstData, net->layers[layer]->output_dim, imgdim, imgdim);
 
     //cv::imwrite("viz/layer" + std::to_string(layer) + ".png", viz);
     //cv::imshow("layer", viz);

tests/shelfnet/shelfnet_mapillary.cpp

@@ -85,6 +85,7 @@ int main()
 
     downloadWeightsifDoNotExist(input_bin, "shelfnet_mapillary", "https://cloud.hipert.unimore.it/s/6WnZCKLjik7xrny/download");
 
+    // Mapillary Vistas has originally 66 classes, but we reduced them to 15 to improve the results on the categories of our interest.
     int classes = 15;
 
     // Network layout