viz yolo3 preprocess

This commit is contained in:
Francesco Gatti
2020-06-16 13:19:30 +02:00
parent 07193fc343
commit 1dfc69ba89
+26 -10
View File
@@ -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;