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

merge

Browse Source
This commit is contained in:
Francesco Gatti
2020-06-01 15:34:59 +02:00
parent 18794d52c2 0548a662f5
commit d4e0d07e09
20 changed files with 1466 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/CenternetDetection.h
+3 -3
View File
@@ -73,9 +73,9 @@ public:
CenternetDetection() {};
~CenternetDetection() {};
bool init(const std::string& tensor_path, const int n_classes=80);
void preprocess(cv::Mat &frame);
void postprocess();
bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1);
void preprocess(cv::Mat &frame, const int bi=0);
void postprocess(const int bi=0,const bool mAP=false);
};
include/tkDNN/DetectionNN.h
+53 -34
View File
@@ -14,7 +14,7 @@
#include "tkdnn.h"
//#define OPENCV_CUDACONTRIB //if OPENCV has been compiled with CUDA and contrib.
// #define OPENCV_CUDACONTRIB //if OPENCV has been compiled with CUDA and contrib.
#ifdef OPENCV_CUDACONTRIB
#include <opencv2/cudawarping.hpp>
@@ -30,10 +30,12 @@ class DetectionNN {
tk::dnn::NetworkRT *netRT = nullptr;
dnnType *input_d;
cv::Size originalSize;
std::vector<cv::Size> originalSize;
cv::Scalar colors[256];
int nBatches = 1;
#ifdef OPENCV_CUDACONTRIB
cv::cuda::GpuMat bgr[3];
cv::cuda::GpuMat imagePreproc;
@@ -47,21 +49,26 @@ class DetectionNN {
* 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) = 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.
*
* @param bi batch index
* @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() = 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*/
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;
@@ -74,49 +81,60 @@ class DetectionNN {
*
* @param tensor_path path to the rt file og the NN.
* @param n_classes number of classes for the given dataset.
* @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) = 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 frame frame to run detection on.
* @param frames frames to run detection on.
* @param cur_batches number of batches to use in inference
* @param save_times if set to true, preprocess, inference and postprocess times
* are saved on a csv file, otherwise not.
* @param times pointer to the output stream where to write times
* @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(cv::Mat &frame, bool save_times=false, std::ofstream *times=nullptr){
if(!frame.data)
FatalError("No image data feed to detection");
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 = frame.size();
printCenteredTitle(" TENSORRT detection ", '=', 30);
originalSize.clear();
if(TKDNN_VERBOSE) printCenteredTitle(" TENSORRT detection ", '=', 30);
{
TIMER_START
preprocess(frame);
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);
}
TIMER_STOP
if(save_times) *times<<t_ns<<";";
}
//do inference
tk::dnn::dataDim_t dim = netRT->input_dim;
dim.n = cur_batches;
{
dim.print();
if(TKDNN_VERBOSE) dim.print();
TIMER_START
netRT->infer(dim, input_d);
TIMER_STOP
dim.print();
if(TKDNN_VERBOSE) dim.print();
stats.push_back(t_ns);
if(save_times) *times<<t_ns<<";";
}
batchDetected.clear();
{
TIMER_START
postprocess();
for(int bi=0; bi<cur_batches;++bi)
postprocess(bi, mAP);
TIMER_STOP
if(save_times) *times<<t_ns<<"\n";
}
@@ -125,10 +143,9 @@ class DetectionNN {
/**
* Method to draw boundixg boxes and labels on a frame.
*
* @param frame orginal frame to draw bounding box on.
* @return frame with boundig boxes.
* @param frames orginal frame to draw bounding box on.
*/
cv::Mat draw(cv::Mat &frame) {
void draw(std::vector<cv::Mat>& frames) {
tk::dnn::box b;
int x0, w, x1, y0, h, y1;
int objClass;
@@ -137,24 +154,26 @@ class DetectionNN {
int baseline = 0;
float font_scale = 0.5;
int thickness = 2;
// draw dets
for(int i=0; i<detected.size(); i++) {
b = detected[i];
x0 = b.x;
x1 = b.x + b.w;
y0 = b.y;
y1 = b.y + b.h;
det_class = classesNames[b.cl];
// draw rectangle
cv::rectangle(frame, 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(frame, cv::Point(x0, y0), cv::Point((x0 + text_size.width - 2), (y0 - text_size.height - 2)), colors[b.cl], -1);
cv::putText(frame, det_class, cv::Point(x0, (y0 - (baseline / 2))), cv::FONT_HERSHEY_SIMPLEX, font_scale, cv::Scalar(255, 255, 255), thickness);
// 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);
}
}
return frame;
}
};
include/tkDNN/Layer.h
+2 -1
View File
@@ -535,6 +535,7 @@ struct box {
int cl;
float x, y, w, h;
float prob;
std::vector<float> probs;
void print()
{
@@ -581,7 +582,7 @@ public:
dnnType *predictions;
static const int MAX_DETECTIONS = 2048;
static const int MAX_DETECTIONS = 8192;
static Yolo::detection *allocateDetections(int nboxes, int classes);
static void mergeDetections(Yolo::detection *dets, int ndets, int classes);
};
include/tkDNN/MobilenetDetection.h
+3 -3
View File
@@ -65,9 +65,9 @@ public:
MobilenetDetection() {};
~MobilenetDetection() {};
bool init(const std::string& tensor_path, const int n_classes);
void preprocess(cv::Mat &frame);
void postprocess();
bool init(const std::string& tensor_path, const int n_classes, const int n_batches=1);
void preprocess(cv::Mat &frame, const int bi=0);
void postprocess(const int bi=0,const bool mAP=false);
};
include/tkDNN/Yolo3Detection.h
+3 -3
View File
@@ -24,9 +24,9 @@ public:
Yolo3Detection() {};
~Yolo3Detection() {};
bool init(const std::string& tensor_path, const int n_classes=80);
void preprocess(cv::Mat &frame);
void postprocess();
bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1);
void preprocess(cv::Mat &frame, const int bi=0);
void postprocess(const int bi=0,const bool mAP=false);
};
include/tkDNN/evaluation.h
+3
View File
@@ -108,6 +108,9 @@ void computeTPFPFN( std::vector<Frame> &images,const int classes,
bool verbose=false, const bool write_on_file=false,
std::string net="");
void printJsonCOCOFormat(std::ofstream *out_file, const std::string image_path, std::vector<tk::dnn::box> bbox, const int classes, const int w, const int h);
}}
#endif /*EVALUATION_H*/
include/tkDNN/utils.h
+5 -3
View File
@@ -36,16 +36,18 @@
#define COL_PURPLEB "\033[1;35m"
#define COL_CYANB "\033[1;36m"
#define TKDNN_VERBOSE 0
// Simple Timer
#define TIMER_START timespec start, end; \
clock_gettime(CLOCK_MONOTONIC, &start);
#define TIMER_STOP_C(col) clock_gettime(CLOCK_MONOTONIC, &end); \
#define TIMER_STOP_C(col, show) clock_gettime(CLOCK_MONOTONIC, &end); \
double t_ns = ((double)(end.tv_sec - start.tv_sec) * 1.0e9 + \
(double)(end.tv_nsec - start.tv_nsec))/1.0e6; \
std::cout<<col<<"Time:"<<std::setw(16)<<t_ns<<" ms\n"<<COL_END;
if(show) std::cout<<col<<"Time:"<<std::setw(16)<<t_ns<<" ms\n"<<COL_END;
#define TIMER_STOP TIMER_STOP_C(COL_CYANB)
#define TIMER_STOP TIMER_STOP_C(COL_CYANB, TKDNN_VERBOSE)
/********************************************************
* Prints the error message, and exits