diff --git a/demo/demo/seg_demo.cpp b/demo/demo/seg_demo.cpp index bd31a1a..482d1ac 100644 --- a/demo/demo/seg_demo.cpp +++ b/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 - if(n_classes == 15 && net_name == "shelfnet_mapillary_fp32") - mapillary_15 = true; + bool mapillary_15_colormap = false; + if( n_classes == 15 && + ( 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(); diff --git a/include/tkDNN/NetworkViz.h b/include/tkDNN/NetworkViz.h index 2e2c3f4..514fefc 100644 --- a/include/tkDNN/NetworkViz.h +++ b/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); }} diff --git a/include/tkDNN/SegmentationNN.h b/include/tkDNN/SegmentationNN.h index ec6ac79..d50bbc4 100644 --- a/include/tkDNN/SegmentationNN.h +++ b/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); diff --git a/README_seg.md b/readme/README_seg.md similarity index 76% rename from README_seg.md rename to readme/README_seg.md index 6717f0d..e257492 100644 --- a/README_seg.md +++ b/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 -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) -./demo shelfnet_mapillary_fp32.rt ../demo/yolo_test.mp4 1 15 +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 +./test_shelfnet # run the yolo test (is slow) +./demo shelfnet_fp32.rt ../demo/yolo_test.mp4 1 19 ``` In general the demo program takes the following parameters: ``` @@ -37,11 +37,13 @@ where * `````` if set to 0 the demo will not resize the input frames, but use it as it is, otherwise it will resize it. * `````` is `````` is set to 1, then the input frames will be proportionally resized using `````` as width baseline. * `````` if set to 0 the demo will not show the visualization but save the video into result.mp4 (if n-batches ==1) -* `````` 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. +* `````` 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 - + +![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_realtime1 | [Cityscapes](https://www.cityscapes-dataset.com/) | 19 | 1024x1024 | [weights](https://cloud.hipert.unimore.it/s/mEDZMRJaGCFWSJF/download) | | shelfnet_berkeley | ShelfNet18_realtime1 | [DeepDrive](https://bdd-data.berkeley.edu/) | 20 | 1024x1024 | [weights](https://cloud.hipert.unimore.it/s/m92e7QdD9gYMF7f/download) | -| shelfnet_mapillary | ShelfNet18_realtime1 | [Mapillary Vistas](https://www.mapillary.com/dataset/vistas?pKey=aFWuj_m4nGoq3-tDz5KAqQ)* | 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. \ No newline at end of file + +## 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. \ No newline at end of file diff --git a/readme/output.gif b/readme/output.gif new file mode 100644 index 0000000..64a9396 Binary files /dev/null and b/readme/output.gif differ diff --git a/src/NetworkViz.cpp b/src/NetworkViz.cpp index 4a3f113..aaf5d55 100644 --- a/src/NetworkViz.cpp +++ b/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(ar*imgdim, imgdim); + cv::Size vdim(img_w, img_h); 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); diff --git a/tests/shelfnet/shelfnet_mapillary.cpp b/tests/shelfnet/shelfnet_mapillary.cpp index 1cb003b..c52d699 100644 --- a/tests/shelfnet/shelfnet_mapillary.cpp +++ b/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