From ca639958aa34b459b3b17055e7ddc99ff1b769d7 Mon Sep 17 00:00:00 2001
From: Biparnak Roy <54219771+biparnakroy@users.noreply.github.com>
Date: Sat, 25 Jul 2020 21:05:44 +0530
Subject: [PATCH] added YOLO output

---
 include/tkDNN/DetectionNN.h | 232 +++++++++++++++++++++---------------
 1 file changed, 133 insertions(+), 99 deletions(-)

diff --git a/include/tkDNN/DetectionNN.h b/include/tkDNN/DetectionNN.h
index 030cf8f..c70247d 100644
--- a/include/tkDNN/DetectionNN.h
+++ b/include/tkDNN/DetectionNN.h
@@ -3,11 +3,13 @@
 
 #include <iostream>
 #include <signal.h>
-#include <stdlib.h>    
+#include <stdlib.h>
 #include <unistd.h>
 #include <mutex>
 #include "utils.h"
 
+#include <iomanip>
+
 #include <opencv2/core/core.hpp>
 #include <opencv2/highgui/highgui.hpp>
 #include <opencv2/imgproc/imgproc.hpp>
@@ -21,39 +23,42 @@
 #include <opencv2/cudaarithm.hpp>
 #endif
 
+namespace tk
+{
+    namespace dnn
+    {
 
-namespace tk { namespace dnn {
+        class DetectionNN
+        {
 
-class DetectionNN {
+        protected:
+            tk::dnn::NetworkRT *netRT = nullptr;
+            dnnType *input_d;
 
-    protected:
-        tk::dnn::NetworkRT *netRT = nullptr;
-        dnnType *input_d;
+            std::vector<cv::Size> originalSize;
 
-        std::vector<cv::Size> originalSize;
+            cv::Scalar colors[256];
 
-        cv::Scalar colors[256];
-
-        int nBatches = 1;
+            int nBatches = 1;
 
 #ifdef OPENCV_CUDACONTRIB
-        cv::cuda::GpuMat bgr[3];
-        cv::cuda::GpuMat imagePreproc;
+            cv::cuda::GpuMat bgr[3];
+            cv::cuda::GpuMat imagePreproc;
 #else
-        cv::Mat bgr[3];
-        cv::Mat imagePreproc;
-        dnnType *input;
+            cv::Mat bgr[3];
+            cv::Mat imagePreproc;
+            dnnType *input;
 #endif
 
-        /**
+            /**
          * 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, 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 
          * boundig boxes. 
          * 
@@ -61,21 +66,21 @@ 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
          */
-        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:
-        int classes = 0;
-        float confThreshold = 0.3; /*threshold on the confidence of the boxes*/
+        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;
+            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;
 
-        DetectionNN() {};
-        ~DetectionNN(){};
+            DetectionNN(){};
+            ~DetectionNN(){};
 
-        /**
+            /**
          * Method used to inialize the class, allocate memory and compute 
          * needed data.
          * 
@@ -84,9 +89,9 @@ class DetectionNN {
          * @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) = 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 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 
          *            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){
-            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); 
+            void update(std::vector<cv::Mat> &frames, const int cur_batches = 1, bool save_times = false, std::ofstream *times = nullptr, const bool mAP = false)
             {
-                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);    
+                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
+                    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.
          * 
          * @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) {
-            tk::dnn::box b;
-            int x0, w, x1, y0, h, y1;
-            int objClass;
-            std::string det_class;
+            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;
 
-            int baseline = 0;
-            float font_scale = 0.5;
-            int thickness = 2;   
+                //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;
 
-            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];
+                int baseline = 0;
+                float font_scale = 0.5;
+                int thickness = 2;
 
-                    // draw rectangle
-                    cv::rectangle(frames[bi], 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(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);
+                        //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);
+
+                        // 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*/