From 98173ec15c9a3b54c4cace8b1142454e41f8b9d7 Mon Sep 17 00:00:00 2001
From: Micaela Verucchi <micaelaverucchi@gmail.com>
Date: Sat, 30 May 2020 17:57:48 +0200
Subject: [PATCH] Add darknetParseFields

Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>
---
 include/tkDNN/DarknetParser.h | 41 ++++++++++++++++++++++++++++++++++-
 tests/yolo3/yolo3.cpp         |  5 +++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/include/tkDNN/DarknetParser.h b/include/tkDNN/DarknetParser.h
index c921d2c..09d8dce 100644
--- a/include/tkDNN/DarknetParser.h
+++ b/include/tkDNN/DarknetParser.h
@@ -7,9 +7,21 @@ namespace tk { namespace dnn {
     struct darknetFields_t{
         int width = 0;
         int height = 0;
+        int channels = 0;
+        int batch_normalize=0;
+        int filters=0;
+        int size=0;
+        int stride=0;
+        int pad=0;
+        std::string activation = "";
 
     };
 
+    std::ostream& operator<<(std::ostream& os, const darknetFields_t& f){
+        os << f.width << " " << f.height << " " << f.channels << " " << f.batch_normalize<< " " << f.filters<< " " << f.size<< " " << f.stride << " " << f.pad << " " << f.activation;
+        return os;
+    }
+
     std::string darknetParseType(const std::string& line){
             size_t start = line.find("[");
             size_t end  = line.find("]");
@@ -20,8 +32,35 @@ namespace tk { namespace dnn {
             return type;
     }
 
-    darknetFields_t parseFields(const std::string& line){
+    bool divideNameAndValue(const std::string& line, std::string&name, std::string& value){
+        size_t sep = line.find("=");
+        if(sep == std::string::npos)
+            return false;
         
+        name   = line.substr(0, sep);   
+        value  = line.substr(sep+1, line.size() - (sep+1));   
+        return true;
+    }
+
+    bool darknetParseFields(const std::string& line, darknetFields_t& fields){
+
+        std::string name,value;
+        if(!divideNameAndValue(line, name, value))
+            return false;
+        std::cout<<name<<std::endl;
+        std::cout<<value<<std::endl;
+        if(name == "width")
+            fields.width = std::stoi(value);
+        else if (name == "height")
+            fields.height = std::stoi(value);
+        else if (name == "channels")
+            fields.channels = std::stoi(value);
+        else if (name == "batch_normalize")
+            fields.batch_normalize = std::stoi(value);
+        else if (name == "filters")
+            fields.filters = std::stoi(value);
+        
+        return true;
     }
 
 
diff --git a/tests/yolo3/yolo3.cpp b/tests/yolo3/yolo3.cpp
index c27a90a..f74faa5 100644
--- a/tests/yolo3/yolo3.cpp
+++ b/tests/yolo3/yolo3.cpp
@@ -7,6 +7,11 @@ int main() {
 
     tk::dnn::Network *net = tk::dnn::DarknetParser("../tests/yolo3/yolov3.cfg");
 
+    tk::dnn::darknetFields_t f;
+    tk::dnn::darknetParseFields("width=40", f);
+    tk::dnn::darknetParseFields("height=40", f);
+    std::cout<<f<<std::endl;
+
     // Network layout
     //tk::dnn::dataDim_t dim(1, 3, 416, 416, 1);
     //tk::dnn::Network net(dim);