added YOLO output

include/tkDNN/DetectionNN.h

@@ -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*/