From 296a6cbc87912a12c662090a39d3d6522b65b85a Mon Sep 17 00:00:00 2001
From: Davide Sapienza <sapienza.dav@gmail.com>
Date: Tue, 3 Mar 2020 17:08:25 +0100
Subject: [PATCH] Add draw method to Yolo3Detection class.

Update yolo demo

Signed-off-by: Davide Sapienza <sapienza.dav@gmail.com>
---
 demo/demo/demo_yolo3.cpp       | 23 +----------------------
 include/tkDNN/Yolo3Detection.h |  2 +-
 src/Yolo3Detection.cpp         | 31 +++++++++++++++++++++++++++++++
 3 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/demo/demo/demo_yolo3.cpp b/demo/demo/demo_yolo3.cpp
index 2b34386..0abcc2c 100644
--- a/demo/demo/demo_yolo3.cpp
+++ b/demo/demo/demo_yolo3.cpp
@@ -67,28 +67,7 @@ int main(int argc, char *argv[]) {
         // TODO: async infer
         yolo.update(dnn_input);
 
-        // draw dets
-        for(int i=0; i<yolo.detected.size(); i++) {
-            tk::dnn::box b = yolo.detected[i];
-            int x0   				= b.x;
-            int x1   				= b.x + b.w;
-            int y0   				= b.y;
-            int y1   				= b.y + b.h;
-            std::string det_class 	= yolo.getYoloLayer()->classesNames[b.cl];
-            float prob 				= b.prob;
-
-            // std::cout<<det_class<<" ("<<prob<<"): "<<x0<<" "<<y0<<" "<<x1<<" "<<y1<<"\n";
-			// draw rectangle
-            cv::rectangle(frame, cv::Point(x0, y0), cv::Point(x1, y1), yolo.colors[b.cl], 2); 
-
-	        // draw label
-            int baseline = 0;
-            float fontScale = 0.5;
-            int thickness = 2;
-            cv::Size textSize = getTextSize(det_class, cv::FONT_HERSHEY_SIMPLEX, fontScale, thickness, &baseline);
-            cv::rectangle(frame, cv::Point(x0, y0), cv::Point((x0 + textSize.width - 2), (y0 - textSize.height - 2)), yolo.colors[b.cl], -1);                      
-            cv::putText(frame, det_class, cv::Point(x0, (y0 - (baseline / 2))), cv::FONT_HERSHEY_SIMPLEX, fontScale, cv::Scalar(255, 255, 255), thickness);
-        }
+        frame = yolo.draw(frame);
         
         cv::imshow("detection", frame);
         cv::waitKey(1);
diff --git a/include/tkDNN/Yolo3Detection.h b/include/tkDNN/Yolo3Detection.h
index 72ea3f8..3700fa6 100644
--- a/include/tkDNN/Yolo3Detection.h
+++ b/include/tkDNN/Yolo3Detection.h
@@ -56,7 +56,7 @@ class Yolo3Detection {
          * @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) {
diff --git a/src/Yolo3Detection.cpp b/src/Yolo3Detection.cpp
index f21efa7..08742c6 100644
--- a/src/Yolo3Detection.cpp
+++ b/src/Yolo3Detection.cpp
@@ -60,6 +60,37 @@ bool Yolo3Detection::init(std::string tensor_path) {
     return true;
 } 
 
+cv::Mat Yolo3Detection::draw(cv::Mat &imageORIG) {
+
+    tk::dnn::box b;
+    int x0, w, x1, y0, h, y1;
+    int objClass;
+    std::string det_class;
+    float prob;
+    int baseline = 0;
+    float fontScale = 0.5;
+    int thickness = 2;   
+    // draw dets
+    for(int i=0; i<detected.size(); i++) {
+        b           = detected[i];
+        x0   		= b.x;
+        x1   		= b.x + b.w;
+        y0   		= b.y;
+        y1   		= b.y + b.h;
+        det_class 	= getYoloLayer()->classesNames[b.cl];
+        prob    	= b.prob;
+
+        // std::cout<<det_class<<" ("<<prob<<"): "<<x0<<" "<<y0<<" "<<x1<<" "<<y1<<"\n";
+        // draw rectangle
+        cv::rectangle(imageORIG, cv::Point(x0, y0), cv::Point(x1, y1), colors[b.cl], 2); 
+
+        // draw label
+        cv::Size textSize = getTextSize(det_class, cv::FONT_HERSHEY_SIMPLEX, fontScale, thickness, &baseline);
+        cv::rectangle(imageORIG, cv::Point(x0, y0), cv::Point((x0 + textSize.width - 2), (y0 - textSize.height - 2)), colors[b.cl], -1);                      
+        cv::putText(imageORIG, det_class, cv::Point(x0, (y0 - (baseline / 2))), cv::FONT_HERSHEY_SIMPLEX, fontScale, cv::Scalar(255, 255, 255), thickness);
+    }
+    return imageORIG;
+}
 
 void Yolo3Detection::update(cv::Mat &imageORIG) {