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 <iostream>
#include <signal.h> #include <signal.h>
#include <stdlib.h> #include <stdlib.h>
#include <unistd.h> #include <unistd.h>
#include <mutex> #include <mutex>
#include "utils.h" #include "utils.h"
#include <iomanip>
#include <opencv2/core/core.hpp> #include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp> #include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp> #include <opencv2/imgproc/imgproc.hpp>
@@ -21,39 +23,42 @@
#include <opencv2/cudaarithm.hpp> #include <opencv2/cudaarithm.hpp>
#endif #endif
namespace tk
{
namespace dnn
{
namespace tk { namespace dnn { class DetectionNN
{
class DetectionNN { protected:
tk::dnn::NetworkRT *netRT = nullptr;
dnnType *input_d;
protected: std::vector<cv::Size> originalSize;
tk::dnn::NetworkRT *netRT = nullptr;
dnnType *input_d;
std::vector<cv::Size> originalSize; cv::Scalar colors[256];
cv::Scalar colors[256]; int nBatches = 1;
int nBatches = 1;
#ifdef OPENCV_CUDACONTRIB #ifdef OPENCV_CUDACONTRIB
cv::cuda::GpuMat bgr[3]; cv::cuda::GpuMat bgr[3];
cv::cuda::GpuMat imagePreproc; cv::cuda::GpuMat imagePreproc;
#else #else
cv::Mat bgr[3]; cv::Mat bgr[3];
cv::Mat imagePreproc; cv::Mat imagePreproc;
dnnType *input; dnnType *input;
#endif #endif
/** /**
* This method preprocess the image, before feeding it to the NN. * This method preprocess the image, before feeding it to the NN.
* *
* @param frame original frame to adapt for inference. * @param frame original frame to adapt for inference.
* @param bi batch index * @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 * This method postprocess the output of the NN to obtain the correct
* boundig boxes. * boundig boxes.
* *
@@ -61,21 +66,21 @@ class DetectionNN {
* @param mAP set to true only if all the probabilities for a bounding * @param mAP set to true only if all the probabilities for a bounding
* box are needed, as in some cases for the mAP calculation * 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: public:
int classes = 0; int classes = 0;
float confThreshold = 0.3; /*threshold on the confidence of the boxes*/ float confThreshold = 0.3; /*threshold on the confidence of the boxes*/
std::vector<tk::dnn::box> detected; /*bounding boxes in output*/ std::vector<tk::dnn::box> detected; /*bounding boxes in output*/
std::vector<std::vector<tk::dnn::box>> batchDetected; /*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<double> stats; /*keeps track of inference times (ms)*/
std::vector<std::string> classesNames; std::vector<std::string> classesNames;
DetectionNN() {}; DetectionNN(){};
~DetectionNN(){}; ~DetectionNN(){};
/** /**
* Method used to inialize the class, allocate memory and compute * Method used to inialize the class, allocate memory and compute
* needed data. * needed data.
* *
@@ -84,9 +89,9 @@ class DetectionNN {
* @param n_batches maximum number of batches to use in inference * @param n_batches maximum number of batches to use in inference
* @return true if everything is correct, false otherwise. * @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. * This method performs the whole detection of the NN.
* *
* @param frames frames to run detection on. * @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 * @param mAP set to true only if all the probabilities for a bounding
* box are needed, as in some cases for the mAP calculation * 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){ 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);
{ {
TKDNN_TSTART if (save_times && times == nullptr)
for(int bi=0; bi<cur_batches;++bi){ FatalError("save_times set to true, but no valid ofstream given");
if(!frames[bi].data) if (cur_batches > nBatches)
FatalError("No image data feed to detection"); FatalError("A batch size greater than nBatches cannot be used");
originalSize.push_back(frames[bi].size());
preprocess(frames[bi], bi); 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. * Method to draw boundixg boxes and labels on a frame.
* *
* @param frames orginal frame to draw bounding box on. * @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) { void draw(std::vector<cv::Mat> &frames, bool ext_yolo)
tk::dnn::box b; {
int x0, w, x1, y0, h, y1; tk::dnn::box b;
int objClass; int x0, w, x1, y0, h, y1;
std::string det_class; int objClass;
std::string det_class;
int baseline = 0; //yolo detctions output
float font_scale = 0.5; std::string yoloBox;
int thickness = 2; 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){ int baseline = 0;
// draw dets float font_scale = 0.5;
for(int i=0; i<batchDetected[bi].size(); i++) { int thickness = 2;
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 rectangle for (int bi = 0; bi < frames.size(); ++bi)
cv::rectangle(frames[bi], cv::Point(x0, y0), cv::Point(x1, y1), colors[b.cl], 2); {
// 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 //yolo stuff
cv::Size text_size = getTextSize(det_class, cv::FONT_HERSHEY_SIMPLEX, font_scale, thickness, &baseline); if (ext_yolo)
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); 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*/ #endif /* DETECTIONNN_H*/