Francesco Gatti
140 lines
4.4 KiB
C++
140 lines
4.4 KiB
C++
#include "Yolo3Detection.h"
|
|
|
|
namespace tk { namespace dnn {
|
|
|
|
bool Yolo3Detection::init(std::string tensor_path) {
|
|
|
|
//const char *tensor_path = "../data/yolo3/yolo3_berkeley.rt";
|
|
|
|
//convert network to tensorRT
|
|
std::cout<<(tensor_path).c_str()<<"\n";
|
|
netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() );
|
|
|
|
if(netRT->pluginFactory->n_yolos != 3) {
|
|
FatalError("this is not yolo3");
|
|
}
|
|
|
|
for(int i=0; i<netRT->pluginFactory->n_yolos; i++) {
|
|
YoloRT *yRT = netRT->pluginFactory->yolos[i];
|
|
classes = yRT->classes;
|
|
num = yRT->num;
|
|
|
|
// make a yolo layer for interpret predictions
|
|
yolo[i] = new tk::dnn::Yolo(nullptr, classes, num, nullptr); // yolo without input and bias
|
|
memcpy(yolo[i]->mask_h, yRT->mask, sizeof(dnnType)*num);
|
|
memcpy(yolo[i]->bias_h, yRT->bias, sizeof(dnnType)*num*3*2);
|
|
yolo[i]->input_dim = yolo[i]->output_dim = tk::dnn::dataDim_t(1, yRT->c, yRT->h, yRT->w);
|
|
}
|
|
|
|
dets = tk::dnn::Yolo::allocateDetections(tk::dnn::Yolo::MAX_DETECTIONS, classes);
|
|
|
|
checkCuda(cudaMallocHost(&input, sizeof(dnnType)*netRT->input_dim.tot()));
|
|
checkCuda(cudaMalloc(&input_d, sizeof(dnnType)*netRT->input_dim.tot()));
|
|
|
|
|
|
// class colors precompute
|
|
for(int c=0; c<classes; c++) {
|
|
int cc = c+1;
|
|
double d = 1.0*( (cc%16)/8 );
|
|
double r = 1.0*( (cc%8)/4 ) + (0.5*d);
|
|
double g = 1.0*( (cc%4)/2 ) + (0.5*d);
|
|
double b = 1.0*( (cc%2)/1 ) + (0.5*d);
|
|
if(r > 1) r = 1;
|
|
if(g > 1) g = 1;
|
|
if(b > 1) b = 1;
|
|
//std::cout<<r<<" "<<g<<" "<<b<<"\n";
|
|
colors[c] = cv::Scalar(int(255.0*b), int(255.0*g), int(255.0*r));
|
|
}
|
|
return true;
|
|
}
|
|
|
|
|
|
void Yolo3Detection::update(cv::Mat &imageORIG) {
|
|
|
|
if(!imageORIG.data) {
|
|
std::cout<<"YOLO: NO IMAGE DATA\n";
|
|
return;
|
|
}
|
|
float xRatio = float(imageORIG.cols) / float(netRT->input_dim.w);
|
|
float yRatio = float(imageORIG.rows) / float(netRT->input_dim.h);
|
|
|
|
resize(imageORIG, imageORIG, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
|
|
imageORIG.convertTo(imageF, CV_32FC3, 1/255.0);
|
|
|
|
//split channels
|
|
cv::split(imageF,bgr);//split source
|
|
|
|
//write channels
|
|
int idx = 0;
|
|
memcpy((void*)&input[idx], (void*)bgr[2].data, imageF.rows*imageF.cols*sizeof(dnnType));
|
|
idx = imageF.rows*imageF.cols;
|
|
memcpy((void*)&input[idx], (void*)bgr[1].data, imageF.rows*imageF.cols*sizeof(dnnType));
|
|
idx *= 2;
|
|
memcpy((void*)&input[idx], (void*)bgr[0].data, imageF.rows*imageF.cols*sizeof(dnnType));
|
|
|
|
//DO INFERENCE
|
|
dnnType *rt_out[3];
|
|
tk::dnn::dataDim_t dim = netRT->input_dim;
|
|
checkCuda(cudaMemcpyAsync(input_d, input, dim.tot()*sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream));
|
|
|
|
printCenteredTitle(" TENSORRT inference ", '=', 30); {
|
|
dim.print();
|
|
TIMER_START
|
|
netRT->infer(dim, input_d);
|
|
TIMER_STOP
|
|
dim.print();
|
|
}
|
|
|
|
TIMER_START
|
|
// compute dets
|
|
ndets = 0;
|
|
for(int i=0; i<3; i++) {
|
|
rt_out[i] = (dnnType*)netRT->buffersRT[i+1];
|
|
yolo[i]->dstData = rt_out[i];
|
|
yolo[i]->computeDetections(dets, ndets, netRT->input_dim.w, netRT->input_dim.h, thresh);
|
|
}
|
|
tk::dnn::Yolo::mergeDetections(dets, ndets, classes);
|
|
TIMER_STOP
|
|
|
|
// fill detected
|
|
detected.clear();
|
|
for(int j=0; j<ndets; j++) {
|
|
tk::dnn::Yolo::box b = dets[j].bbox;
|
|
int x0 = (b.x-b.w/2.);
|
|
int x1 = (b.x+b.w/2.);
|
|
int y0 = (b.y-b.h/2.);
|
|
int y1 = (b.y+b.h/2.);
|
|
int obj_class = -1;
|
|
float prob = 0;
|
|
for(int c=0; c<classes; c++) {
|
|
if(dets[j].prob[c] >= thresh) {
|
|
obj_class = c;
|
|
prob = dets[j].prob[c];
|
|
}
|
|
}
|
|
|
|
if(obj_class >= 0) {
|
|
//std::cout<<obj_class<<" ("<<prob<<"): "<<x0<<" "<<y0<<" "<<x1<<" "<<y1<<"\n";
|
|
//cv::rectangle(image, cv::Point(x0, y0), cv::Point(x1, y1), colors[obj_class], 2);
|
|
|
|
// convert to image coords
|
|
x0 = xRatio*x0;
|
|
x1 = xRatio*x1;
|
|
y0 = yRatio*y0;
|
|
y1 = yRatio*y1;
|
|
|
|
tk::dnn::box res;
|
|
res.cl = obj_class;
|
|
res.prob = prob;
|
|
res.x = x0;
|
|
res.y = y0;
|
|
res.w = x1 - x0;
|
|
res.h = y1 - y0;
|
|
detected.push_back(res);
|
|
}
|
|
}
|
|
|
|
}
|
|
|
|
}}
|