Refactoring for detection NN

Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>
2020-03-20 21:14:12 +01:00
parent 997e857e32
commit 8b9516da97
11 changed files with 637 additions and 678 deletions
@@ -1,73 +1,36 @@
-#ifndef YOLODETECTION_H
-#define YOLODETECTION_H
+#ifndef Yolo3Detection_H
+#define Yolo3Detection_H
+#include <opencv2/videoio.hpp>
+#include "opencv2/opencv.hpp"

-#include <iostream>
-#include <signal.h>
-#include <stdlib.h>     /* srand, rand */
-#include <unistd.h>
-#include <mutex>
-#include "utils.h"
+#include "DetectionNN.h"

-#include <opencv2/core/core.hpp>
-#include <opencv2/highgui/highgui.hpp>
-#include <opencv2/imgproc/imgproc.hpp>
+namespace tk { namespace dnn { 

-#include "tkdnn.h"
+class Yolo3Detection : public DetectionNN
+{
+private:
+    int num = 0;
+    int nMasks = 0;
+    int nDets = 0;
+    tk::dnn::Yolo::detection *dets = nullptr;
+    tk::dnn::Yolo* yolo[3];

-namespace tk { namespace dnn {
-
-/**
- * 
- * @author Francesco Gatti
- */
-class Yolo3Detection {
-
-    private:
-        tk::dnn::NetworkRT *netRT = nullptr;
-        tk::dnn::Yolo* yolo[3];
-        dnnType *input, *input_d;
-
-        int ndets = 0;
-        tk::dnn::Yolo::detection *dets = nullptr;
-
-        cv::Mat imageF;
-        cv::Mat bgr[3]; 
-
-    public:
-        int classes = 0;
-        int num = 0;
-        int n_masks = 0;
-        float thresh = 0.3;
-        cv::Scalar colors[256];
-
-        // this is filled with results
-        std::vector<tk::dnn::box> detected;
-
-        // keep track of inference times (ms)
-        std::vector<double> stats;
-
-        Yolo3Detection() {}
-
-        virtual ~Yolo3Detection() {}
-
-        /**
-         * Method used for inizialize the class
-         * 
-         * @return Success of the initialization
-         */
-        bool init(std::string tensor_path);
-        cv::Mat draw(cv::Mat &frame);
-        void update(cv::Mat &frame);
-
-        tk::dnn::Yolo* getYoloLayer(int n=0) {
-            if(n<3)
-                return yolo[n];
-            else 
-                return nullptr;
-        }
+    tk::dnn::Yolo* getYoloLayer(int n=0);
+    
+public:
+    Yolo3Detection() {};
+    ~Yolo3Detection() {}; 

+    bool init(const std::string& tensor_path, const int n_classes=80);
+    void preprocess(cv::Mat &frame);
+    void update(cv::Mat &frame);
+    void postprocess(dnnType **rt_out, const int n_out);
+    cv::Mat draw(cv::Mat &frame);
 };

-}}

-#endif /* YOLODETECTION_H*/
+} // namespace dnn
+} // namespace tk
+
+#endif /* Yolo3Detection_H*/