mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

added YOLO output

Browse Source
This commit is contained in:
Biparnak Roy
2020-07-25 21:05:44 +05:30
committed by GitHub
parent e976c71edc
commit ca639958aa
1 changed files with 133 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/DetectionNN.h
+133 -99
View File
@@ -3,11 +3,13 @@
#include <iostream>
#include <signal.h>
#include <stdlib.h>
#include <stdlib.h>
#include <unistd.h>
#include <mutex>
#include "utils.h"
#include <iomanip>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
@@ -21,39 +23,42 @@
#include <opencv2/cudaarithm.hpp>
#endif
namespace tk
{
namespace dnn
{
namespace tk { namespace dnn {
class DetectionNN
{
class DetectionNN {
protected:
tk::dnn::NetworkRT *netRT = nullptr;
dnnType *input_d;
protected:
tk::dnn::NetworkRT *netRT = nullptr;
dnnType *input_d;
std::vector<cv::Size> originalSize;
std::vector<cv::Size> originalSize;
cv::Scalar colors[256];
cv::Scalar colors[256];
int nBatches = 1;
int nBatches = 1;
#ifdef OPENCV_CUDACONTRIB
cv::cuda::GpuMat bgr[3];
cv::cuda::GpuMat imagePreproc;
cv::cuda::GpuMat bgr[3];
cv::cuda::GpuMat imagePreproc;
#else
cv::Mat bgr[3];
cv::Mat imagePreproc;
dnnType *input;
cv::Mat bgr[3];
cv::Mat imagePreproc;
dnnType *input;
#endif
/**
/**
* This method preprocess the image, before feeding it to the NN.
*
* @param frame original frame to adapt for inference.
* @param bi batch index
*/
virtual void preprocess(cv::Mat &frame, const int bi=0) = 0;
virtual void preprocess(cv::Mat &frame, const int bi = 0) = 0;
/**
/**
* This method postprocess the output of the NN to obtain the correct
* boundig boxes.
*
@@ -61,21 +66,21 @@ class DetectionNN {
* @param mAP set to true only if all the probabilities for a bounding
* box are needed, as in some cases for the mAP calculation
*/
virtual void postprocess(const int bi=0,const bool mAP=false) = 0;
virtual void postprocess(const int bi = 0, const bool mAP = false) = 0;
public:
int classes = 0;
float confThreshold = 0.3; /*threshold on the confidence of the boxes*/
public:
int classes = 0;
float confThreshold = 0.3; /*threshold on the confidence of the boxes*/
std::vector<tk::dnn::box> detected; /*bounding boxes in output*/
std::vector<std::vector<tk::dnn::box>> batchDetected; /*bounding boxes in output*/
std::vector<double> stats; /*keeps track of inference times (ms)*/
std::vector<std::string> classesNames;
std::vector<tk::dnn::box> detected; /*bounding boxes in output*/
std::vector<std::vector<tk::dnn::box>> batchDetected; /*bounding boxes in output*/
std::vector<double> stats; /*keeps track of inference times (ms)*/
std::vector<std::string> classesNames;
DetectionNN() {};
~DetectionNN(){};
DetectionNN(){};
~DetectionNN(){};
/**
/**
* Method used to inialize the class, allocate memory and compute
* needed data.
*
@@ -84,9 +89,9 @@ class DetectionNN {
* @param n_batches maximum number of batches to use in inference
* @return true if everything is correct, false otherwise.
*/
virtual bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1) = 0;
/**
virtual bool init(const std::string &tensor_path, const int n_classes = 80, const int n_batches = 1) = 0;
/**
* This method performs the whole detection of the NN.
*
* @param frames frames to run detection on.
@@ -97,87 +102,116 @@ class DetectionNN {
* @param mAP set to true only if all the probabilities for a bounding
* box are needed, as in some cases for the mAP calculation
*/
void update(std::vector<cv::Mat>& frames, const int cur_batches=1, bool save_times=false, std::ofstream *times=nullptr, const bool mAP=false){
if(save_times && times==nullptr)
FatalError("save_times set to true, but no valid ofstream given");
if(cur_batches > nBatches)
FatalError("A batch size greater than nBatches cannot be used");
originalSize.clear();
if(TKDNN_VERBOSE) printCenteredTitle(" TENSORRT detection ", '=', 30);
void update(std::vector<cv::Mat> &frames, const int cur_batches = 1, bool save_times = false, std::ofstream *times = nullptr, const bool mAP = false)
{
TKDNN_TSTART
for(int bi=0; bi<cur_batches;++bi){
if(!frames[bi].data)
FatalError("No image data feed to detection");
originalSize.push_back(frames[bi].size());
preprocess(frames[bi], bi);
if (save_times && times == nullptr)
FatalError("save_times set to true, but no valid ofstream given");
if (cur_batches > nBatches)
FatalError("A batch size greater than nBatches cannot be used");
originalSize.clear();
if (TKDNN_VERBOSE)
printCenteredTitle(" TENSORRT detection ", '=', 30);
{
TKDNN_TSTART
for (int bi = 0; bi < cur_batches; ++bi)
{
if (!frames[bi].data)
FatalError("No image data feed to detection");
originalSize.push_back(frames[bi].size());
preprocess(frames[bi], bi);
}
TKDNN_TSTOP
if (save_times)
*times << t_ns << ";";
}
//do inference
tk::dnn::dataDim_t dim = netRT->input_dim;
dim.n = cur_batches;
{
if (TKDNN_VERBOSE)
dim.print();
TKDNN_TSTART
netRT->infer(dim, input_d);
TKDNN_TSTOP
if (TKDNN_VERBOSE)
dim.print();
stats.push_back(t_ns);
if (save_times)
*times << t_ns << ";";
}
batchDetected.clear();
{
TKDNN_TSTART
for (int bi = 0; bi < cur_batches; ++bi)
postprocess(bi, mAP);
TKDNN_TSTOP
if (save_times)
*times << t_ns << "\n";
}
TKDNN_TSTOP
if(save_times) *times<<t_ns<<";";
}
//do inference
tk::dnn::dataDim_t dim = netRT->input_dim;
dim.n = cur_batches;
{
if(TKDNN_VERBOSE) dim.print();
TKDNN_TSTART
netRT->infer(dim, input_d);
TKDNN_TSTOP
if(TKDNN_VERBOSE) dim.print();
stats.push_back(t_ns);
if(save_times) *times<<t_ns<<";";
}
batchDetected.clear();
{
TKDNN_TSTART
for(int bi=0; bi<cur_batches;++bi)
postprocess(bi, mAP);
TKDNN_TSTOP
if(save_times) *times<<t_ns<<"\n";
}
}
/**
/**
* Method to draw boundixg boxes and labels on a frame.
*
* @param frames orginal frame to draw bounding box on.
* @param ext_yolo exports yolo style coorinates of bounding boxes on the terminal
*/
void draw(std::vector<cv::Mat>& frames) {
tk::dnn::box b;
int x0, w, x1, y0, h, y1;
int objClass;
std::string det_class;
void draw(std::vector<cv::Mat> &frames, bool ext_yolo)
{
tk::dnn::box b;
int x0, w, x1, y0, h, y1;
int objClass;
std::string det_class;
int baseline = 0;
float font_scale = 0.5;
int thickness = 2;
//yolo detctions output
std::string yoloBox;
float Yx, Yy, Yw, Yh;
cv::Size sz = frames[0].size();
int imageWidth = sz.width;
int imageHeight = sz.height;
for(int bi=0; bi<frames.size(); ++bi){
// draw dets
for(int i=0; i<batchDetected[bi].size(); i++) {
b = batchDetected[bi][i];
x0 = b.x;
x1 = b.x + b.w;
y0 = b.y;
y1 = b.y + b.h;
det_class = classesNames[b.cl];
int baseline = 0;
float font_scale = 0.5;
int thickness = 2;
// draw rectangle
cv::rectangle(frames[bi], cv::Point(x0, y0), cv::Point(x1, y1), colors[b.cl], 2);
for (int bi = 0; bi < frames.size(); ++bi)
{
// draw dets
for (int i = 0; i < batchDetected[bi].size(); i++)
{
b = batchDetected[bi][i];
x0 = b.x;
x1 = b.x + b.w;
y0 = b.y;
y1 = b.y + b.h;
det_class = classesNames[b.cl];
// draw label
cv::Size text_size = getTextSize(det_class, cv::FONT_HERSHEY_SIMPLEX, font_scale, thickness, &baseline);
cv::rectangle(frames[bi], cv::Point(x0, y0), cv::Point((x0 + text_size.width - 2), (y0 - text_size.height - 2)), colors[b.cl], -1);
cv::putText(frames[bi], det_class, cv::Point(x0, (y0 - (baseline / 2))), cv::FONT_HERSHEY_SIMPLEX, font_scale, cv::Scalar(255, 255, 255), thickness);
//yolo stuff
if (ext_yolo)
{
Yx = (b.x + (int)(b.w / 2)) / imageWidth;
Yy = (b.y + (int)(b.h / 2)) / imageHeight;
Yw = b.w / imageWidth;
Yh = b.h / imageHeight;
std::cout << std::fixed << std::setprecision(6)<<b.cl<<" "<<Yx<<" "<<Yy<<" "<<Yw<<" "<<Yh<<"\n";
}
// draw rectangle
cv::rectangle(frames[bi], cv::Point(x0, y0), cv::Point(x1, y1), colors[b.cl], 2);
// draw label
cv::Size text_size = getTextSize(det_class, cv::FONT_HERSHEY_SIMPLEX, font_scale, thickness, &baseline);
cv::rectangle(frames[bi], cv::Point(x0, y0), cv::Point((x0 + text_size.width - 2), (y0 - text_size.height - 2)), colors[b.cl], -1);
cv::putText(frames[bi], det_class, cv::Point(x0, (y0 - (baseline / 2))), cv::FONT_HERSHEY_SIMPLEX, font_scale, cv::Scalar(255, 255, 255), thickness);
}
}
}
}
};
};
}}
} // namespace dnn
} // namespace tk
#endif /* DETECTIONNN_H*/