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

Enabled cuda pre-process via opencv_contrib

Browse Source
This commit is contained in:
Harijs Grinbergs
2021-11-04 20:39:37 +02:00
parent 18303c06c5
commit e449d344c0
3 changed files with 49 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/DetectionNN.h
+22 -22
View File
@@ -3,10 +3,10 @@
#include <iostream>
#include <signal.h>
#include <stdlib.h>
#include <stdlib.h>
#ifdef __linux__
#include <unistd.h>
#endif
#endif
#include <mutex>
#include "utils.h"
@@ -17,7 +17,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>
@@ -37,7 +37,7 @@ class DetectionNN {
cv::Scalar colors[256];
int nBatches = 1;
int nBatches = 2;
#ifdef OPENCV_CUDACONTRIB
cv::cuda::GpuMat bgr[3];
@@ -57,11 +57,11 @@ class DetectionNN {
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.
*
* 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
* @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;
@@ -79,25 +79,25 @@ class DetectionNN {
~DetectionNN(){};
/**
* Method used to initialize the class, allocate memory and compute
* Method used to initialize the class, allocate memory and compute
* needed data.
*
*
* @param tensor_path path to the rt file of 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, const int n_batches=1, const float conf_thresh=0.3) = 0;
/**
* This method performs the whole detection of the NN.
*
*
* @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
* @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
* @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){
@@ -107,14 +107,14 @@ class DetectionNN {
FatalError("A batch size greater than nBatches cannot be used");
originalSize.clear();
if(TKDNN_VERBOSE) printCenteredTitle(" TENSORRT detection ", '=', 30);
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);
preprocess(frames[bi], bi);
}
TKDNN_TSTOP
if(save_times) *times<<t_ns<<";";
@@ -141,11 +141,11 @@ class DetectionNN {
TKDNN_TSTOP
if(save_times) *times<<t_ns<<"\n";
}
}
}
/**
* Method to draw bounding boxes and labels on a frame.
*
*
* @param frames original frame to draw bounding box on.
*/
void draw(std::vector<cv::Mat>& frames) {
@@ -155,11 +155,11 @@ class DetectionNN {
std::string det_class;
int baseline = 0;
float font_scale = 0.5;
int thickness = 2;
int thickness = 2;
for(int bi=0; bi<frames.size(); ++bi){
// draw dets
for(int i=0; i<batchDetected[bi].size(); i++) {
for(int i=0; i<batchDetected[bi].size(); i++) {
b = batchDetected[bi][i];
x0 = b.x;
x1 = b.x + b.w;
@@ -168,11 +168,11 @@ class DetectionNN {
det_class = classesNames[b.cl];
// draw rectangle
cv::rectangle(frames[bi], cv::Point(x0, y0), cv::Point(x1, y1), colors[b.cl], 2);
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::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);
}
}
include/tkDNN/DetectionNN3D.h
+20 -20
View File
@@ -4,7 +4,7 @@
#include <iostream>
#include <signal.h>
#include <stdlib.h>
#ifdef __linux__
#ifdef __linux__
#include <unistd.h>
#endif
@@ -17,7 +17,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>
@@ -57,11 +57,11 @@ class DetectionNN3D {
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.
*
* 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
* @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;
@@ -69,7 +69,7 @@ class DetectionNN3D {
public:
int classes = 0;
float confThreshold = 0.3; /*threshold on the confidence of the boxes*/
std::vector<tk::dnn::box3D> detected3D; /*bounding boxes in output*/
std::vector<std::vector<tk::dnn::box3D>> batchDetected; /*bounding boxes in output*/
std::vector<double> pre_stats, stats, post_stats, visual_stats; /*keeps track of inference times (ms)*/
@@ -79,29 +79,29 @@ class DetectionNN3D {
~DetectionNN3D(){};
/**
* Method used to initialize the class, allocate memory and compute
* Method used to initialize the class, allocate memory and compute
* needed data.
*
*
* @param tensor_path path to the rt file of 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=3, const int n_batches=1,
virtual bool init(const std::string& tensor_path, const int n_classes=3, const int n_batches=1,
const float conf_thresh=0.3, const std::vector<cv::Mat>& k_calibs=std::vector<cv::Mat>()) = 0;
/**
* This method performs the whole detection of the NN.
*
*
* @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
* @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
* @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,
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");
@@ -109,17 +109,17 @@ class DetectionNN3D {
FatalError("A batch size greater than nBatches cannot be used");
originalSize.clear();
if(TKDNN_VERBOSE) printCenteredTitle(" TENSORRT detection ", '=', 30);
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);
preprocess(frames[bi], bi);
}
TKDNN_TSTOP
pre_stats.push_back(t_ns);
pre_stats.push_back(t_ns);
if(save_times) *times<<t_ns<<";";
}
@@ -149,11 +149,11 @@ class DetectionNN3D {
/**
* Method to draw bounding boxes and labels on a frame.
*
*
* @param frames original frame to draw bounding box on.
*/
*/
virtual void draw(std::vector<cv::Mat>& frames){};
};
}}