From 6d467e5d930d6326ef072887c93e321e13693e54 Mon Sep 17 00:00:00 2001
From: Francesco Gatti <gattifrancesco@hotmail.it>
Date: Tue, 16 Jun 2020 13:19:30 +0200
Subject: [PATCH] viz yolo3 preprocess

---
 tests/darknet/viz_yolo3.cpp | 36 ++++++++++++++++++++++++++----------
 1 file changed, 26 insertions(+), 10 deletions(-)

diff --git a/tests/darknet/viz_yolo3.cpp b/tests/darknet/viz_yolo3.cpp
index 374f7dd..9e53116 100644
--- a/tests/darknet/viz_yolo3.cpp
+++ b/tests/darknet/viz_yolo3.cpp
@@ -1,32 +1,49 @@
 #include<iostream>
 #include<vector>
 #include <opencv2/highgui/highgui.hpp>
+#include <opencv2/imgproc/imgproc.hpp>
 
 #include "tkdnn.h"
 #include "test.h"
 #include "DarknetParser.h"
 #include "NetworkViz.h"
 
-int main() {
+int main(int argc, char *argv[]) {
+    if(argc <2)
+        FatalError("you must provide an input image");
+    std::string input_image = argv[1];
     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");
+    downloadWeightsifDoNotExist(wgs_path, 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
+    // input data
     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; 
+    checkCuda( cudaMalloc(&input_d, sizeof(dnnType)*net->input_dim.tot()));
+
+    // load image
+    cv::Mat frame, frameFloat;
+    frame = cv::imread(input_image);
+    cv::resize(frame, frame, cv::Size(net->input_dim.w, net->input_dim.h));
+    frame.convertTo(frameFloat, CV_32FC3, 1/255.0); 
+
+    //split channels
+    cv::Mat bgr[3];
+    cv::split(frameFloat,bgr);//split source
+
+    //write channels
+    for(int i=0; i<net->input_dim.c; i++) {
+        int idx = i*frameFloat.rows*frameFloat.cols;
+        int ch = net->input_dim.c-1 -i;
+        checkCuda( cudaMemcpy(input_d + idx, (void*)bgr[ch].data, frameFloat.rows*frameFloat.cols*sizeof(dnnType), cudaMemcpyHostToDevice));     
+    }
     
+    tk::dnn::dataDim_t dim =  net->input_dim;     
     dim.print();
     std::cout<<"infer\n";
     net->infer(dim, input_d);    
@@ -44,7 +61,6 @@ int main() {
         //cv::waitKey(0);
     }
 
-    delete [] input_h;
     checkCuda(cudaFree(input_d));
     net->releaseLayers();
     delete net;