From 0f3dbac962abf25660556510356b8984ab95a953 Mon Sep 17 00:00:00 2001
From: Francesco Gatti <gattifrancesco@hotmail.it>
Date: Tue, 16 Jun 2020 12:36:45 +0200
Subject: [PATCH] viz yolo3

---
 include/tkDNN/DarknetParser.h |  2 +-
 include/tkDNN/NetworkViz.h    | 12 ++++++
 include/tkDNN/utils.h         |  4 ++
 src/NetworkViz.cpp            | 69 +++++++++++++++++++++++++++++++++++
 tests/darknet/viz_yolo3.cpp   | 54 +++++++++++++++++++++++++++
 5 files changed, 140 insertions(+), 1 deletion(-)
 create mode 100644 include/tkDNN/NetworkViz.h
 create mode 100644 src/NetworkViz.cpp
 create mode 100644 tests/darknet/viz_yolo3.cpp

diff --git a/include/tkDNN/DarknetParser.h b/include/tkDNN/DarknetParser.h
index cbeed48..cd7fa48 100644
--- a/include/tkDNN/DarknetParser.h
+++ b/include/tkDNN/DarknetParser.h
@@ -1,6 +1,6 @@
 #pragma once
 #include <iostream>
-#include "tkdnn.h"
+#include "tkDNN/tkdnn.h"
 
 namespace tk { namespace dnn {
 
diff --git a/include/tkDNN/NetworkViz.h b/include/tkDNN/NetworkViz.h
new file mode 100644
index 0000000..c8b1bea
--- /dev/null
+++ b/include/tkDNN/NetworkViz.h
@@ -0,0 +1,12 @@
+#pragma once
+#include <iostream>
+#include <opencv2/core/types.hpp>
+#include "tkdnn.h"
+
+namespace tk { namespace dnn {
+
+cv::Mat vizFloat2colorMap(cv::Mat map);
+cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int imgdim);
+cv::Mat vizLayer2Mat(tk::dnn::Network *net, int layer, int imgdim = 1000);
+    
+}}
diff --git a/include/tkDNN/utils.h b/include/tkDNN/utils.h
index aa73e9e..538a3f3 100644
--- a/include/tkDNN/utils.h
+++ b/include/tkDNN/utils.h
@@ -118,4 +118,8 @@ void matrixMulAdd(  cublasHandle_t handle, dnnType* srcData, dnnType* dstData,
 void getMemUsage(double& vm_usage_kb, double& resident_set_kb);
 void printCudaMemUsage();
 void removePathAndExtension(const std::string &full_string, std::string &name);
+static inline bool isCudaPointer(void *data) {
+  cudaPointerAttributes attr;
+  return cudaPointerGetAttributes(&attr, data) == 0;
+}
 #endif //UTILS_H
diff --git a/src/NetworkViz.cpp b/src/NetworkViz.cpp
new file mode 100644
index 0000000..6ac274c
--- /dev/null
+++ b/src/NetworkViz.cpp
@@ -0,0 +1,69 @@
+#include <opencv2/core/core.hpp>
+#include <opencv2/highgui/highgui.hpp>
+#include <opencv2/videoio.hpp>
+#include <opencv2/imgproc/imgproc.hpp>
+#include "tkDNN/NetworkViz.h"
+
+namespace tk { namespace dnn {
+
+cv::Mat vizFloat2colorMap(cv::Mat map) {
+
+    double min;
+    double max;
+    cv::minMaxIdx(map, &min, &max);
+    cv::Mat adjMap;
+    // expand your range to 0..255. Similar to histEq();
+    map.convertTo(adjMap,CV_8UC1, 255 / (max-min), -min); 
+    //return adjMap;
+
+    
+    cv::Mat falseColorsMap;
+    applyColorMap(adjMap, falseColorsMap, cv::COLORMAP_HOT);
+    return falseColorsMap;
+}
+
+cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int imgdim) {
+    dnnType *data = nullptr;
+
+    // copy to CPU
+    if(isCudaPointer(dataInput)) {
+        data = new dnnType[dim.tot()];
+        checkCuda( cudaMemcpy(data, dataInput, dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToHost) );
+    } else {
+        data = dataInput;
+    }
+
+    int gridDim = ceil(sqrt(dim.c));
+    cv::Size gridSize(dim.w*gridDim, dim.h*gridDim);
+    cv::Mat grid = cv::Mat(gridSize, CV_8UC3, cv::Scalar(0));    
+
+    for(int i=0; i<dim.c;i++) {
+        cv::Mat raw = vizFloat2colorMap(cv::Mat(cv::Size(dim.w, dim.h),CV_32FC1, data + dim.w*dim.h*i));
+        int r = i / gridDim;
+        int c = i - r * gridDim;
+        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::Mat viz;
+    cv::resize(grid, viz, vdim, 0, 0, 0);
+    
+    // free memory
+    if(isCudaPointer(dataInput)) {
+        delete [] data;
+    }
+    return viz;
+}
+
+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);
+
+    //cv::imwrite("viz/layer" + std::to_string(layer) + ".png", viz);
+    //cv::imshow("layer", viz);
+    //cv::waitKey(0);
+}
+
+}}
\ No newline at end of file
diff --git a/tests/darknet/viz_yolo3.cpp b/tests/darknet/viz_yolo3.cpp
new file mode 100644
index 0000000..374f7dd
--- /dev/null
+++ b/tests/darknet/viz_yolo3.cpp
@@ -0,0 +1,54 @@
+#include<iostream>
+#include<vector>
+#include <opencv2/highgui/highgui.hpp>
+
+#include "tkdnn.h"
+#include "test.h"
+#include "DarknetParser.h"
+#include "NetworkViz.h"
+
+int main() {
+    std::string bin_path  = "yolo3";
+    std::vector<std::string> input_bins = { 
+        bin_path + "/layers/input.bin"
+    };
+    std::string wgs_path  = bin_path + "/layers";
+    std::string cfg_path  = std::string(TKDNN_PATH) + "/tests/darknet/cfg/yolo3.cfg";
+    std::string name_path = std::string(TKDNN_PATH) + "/tests/darknet/names/coco.names";
+    downloadWeightsifDoNotExist(input_bins[0], bin_path, "https://cloud.hipert.unimore.it/s/jPXmHyptpLoNdNR/download");
+
+    // parse darknet network
+    tk::dnn::Network *net = tk::dnn::darknetParser(cfg_path, wgs_path, name_path);
+    net->print();
+
+    // Load input and infer
+    dnnType *input_d;
+    dnnType *input_h;
+    readBinaryFile(input_bins[0], net->input_dim.tot(), &input_h, &input_d);
+    tk::dnn::dataDim_t dim =  net->input_dim; 
+    
+    dim.print();
+    std::cout<<"infer\n";
+    net->infer(dim, input_d);    
+
+    // output directory
+    std::string output_viz = "viz/";
+    system( (std::string("mkdir -p ") + output_viz).c_str() );
+
+    for(int i=0; i<net->num_layers; i++) {
+        std::string output_png = output_viz + "/layer" + std::to_string(i) + ".png";
+        std::cout<<"saving "<<output_png<<"\n";
+        cv::Mat viz = vizLayer2Mat(net, i);
+        cv::imwrite(output_png, viz);
+        //cv::imshow("layer", viz);
+        //cv::waitKey(0);
+    }
+
+    delete [] input_h;
+    checkCuda(cudaFree(input_d));
+    net->releaseLayers();
+    delete net;
+    return 0;
+}
+
+ 
\ No newline at end of file