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
+50 -16
View File
@@ -8,6 +8,8 @@
#include <mutex>
#include "utils.h"
#include <iomanip>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
@@ -21,10 +23,13 @@
#include <opencv2/cudaarithm.hpp>
#endif
namespace tk
{
namespace dnn
{
namespace tk { namespace dnn {
class DetectionNN {
class DetectionNN
{
protected:
tk::dnn::NetworkRT *netRT = nullptr;
@@ -97,37 +102,44 @@ 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){
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);
if (TKDNN_VERBOSE)
printCenteredTitle(" TENSORRT detection ", '=', 30);
{
TKDNN_TSTART
for(int bi=0; bi<cur_batches;++bi){
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<<";";
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();
if (TKDNN_VERBOSE)
dim.print();
TKDNN_TSTART
netRT->infer(dim, input_d);
TKDNN_TSTOP
if(TKDNN_VERBOSE) dim.print();
if (TKDNN_VERBOSE)
dim.print();
stats.push_back(t_ns);
if(save_times) *times<<t_ns<<";";
if (save_times)
*times << t_ns << ";";
}
batchDetected.clear();
@@ -136,7 +148,8 @@ class DetectionNN {
for (int bi = 0; bi < cur_batches; ++bi)
postprocess(bi, mAP);
TKDNN_TSTOP
if(save_times) *times<<t_ns<<"\n";
if (save_times)
*times << t_ns << "\n";
}
}
@@ -144,20 +157,31 @@ class DetectionNN {
* 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) {
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;
//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;
int baseline = 0;
float font_scale = 0.5;
int thickness = 2;
for(int bi=0; bi<frames.size(); ++bi){
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;
@@ -165,6 +189,16 @@ class DetectionNN {
y1 = b.y + b.h;
det_class = classesNames[b.cl];
//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);
@@ -175,9 +209,9 @@ class DetectionNN {
}
}
}
};
}}
} // namespace dnn
} // namespace tk
#endif /* DETECTIONNN_H*/