From 46c5a080258778407b9bcd18aea0593c5dbbc99b Mon Sep 17 00:00:00 2001
From: Francois Deschenes <francois@tintup.com>
Date: Wed, 25 Jul 2018 23:58:13 -0700
Subject: [PATCH] Added getLayerNames() to Net and improved blobFromImage();

---
 ext/opencv/dnn.cpp     | 19 ++++++++++++++++---
 ext/opencv/dnn_net.cpp | 21 ++++++++++++++++++++-
 ext/opencv/opencv.hpp  |  4 ++++
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/ext/opencv/dnn.cpp b/ext/opencv/dnn.cpp
index 5f69d08..8dbd978 100644
--- a/ext/opencv/dnn.cpp
+++ b/ext/opencv/dnn.cpp
@@ -31,10 +31,23 @@ namespace rubyopencv {
           Check_Type(options, T_HASH);
 
           double scale_factor = NUM2DBL_DEFAULT(HASH_LOOKUP(options, "scale_factor"), 1.0);
-          cv::Size *s = Size::obj2size(HASH_LOOKUP(options, "size"));
-          cv::Scalar *sc = Scalar::obj2scalar(HASH_LOOKUP(options, "mean"));;
+          cv::Size size;
+          cv::Scalar mean;
+          bool swap_rb = RTEST_DEFAULT(HASH_LOOKUP(options, "swap_rb"), true);
+          bool crop = RTEST_DEFAULT(HASH_LOOKUP(options, "crop"), true);
+      	  VALUE tmp = Qnil;
 
-          r = cv::dnn::blobFromImage(*m, scale_factor, *s, *sc);
+      	  tmp = HASH_LOOKUP(options, "size");
+      	  if (!NIL_P(tmp)) {
+      	    size = *(Size::obj2size(tmp));
+      	  }
+
+      	  tmp = HASH_LOOKUP(options, "mean");
+      	  if (!NIL_P(tmp)) {
+      	    mean = *(Scalar::obj2scalar(tmp));
+      	  }
+
+          r = cv::dnn::blobFromImage(*m, scale_factor, size, mean, swap_rb, crop);
         }
 
         b = new cv::Mat(r);
diff --git a/ext/opencv/dnn_net.cpp b/ext/opencv/dnn_net.cpp
index 8e44f25..745c49f 100644
--- a/ext/opencv/dnn_net.cpp
+++ b/ext/opencv/dnn_net.cpp
@@ -96,15 +96,34 @@ namespace rubyopencv {
         return selfptr->empty() ? Qtrue : Qfalse;
       }
 
+      VALUE rb_get_layer_names(VALUE self) {
+        cv::dnn::Net* selfptr = obj2net(self);
+
+        std::vector<cv::String> v = selfptr->getLayerNames();
+        const long size = v.size();
+
+        VALUE layer_names = rb_ary_new_capa(size);
+        for (long i = 0; i < size; i++) {
+          VALUE layer_name = rb_str_new_cstr(v[i].c_str());
+          rb_ary_store(layer_names, i, layer_name);
+        }
+
+        return layer_names;
+      }
+
       void init(VALUE rb_module) {
         rb_klass = rb_define_class_under(rb_module, "Net", rb_cData);
         rb_define_alloc_func(rb_klass, rb_allocate);
 
         rb_define_private_method(rb_klass, "initialize", RUBY_METHOD_FUNC(rb_initialize), 0);
 
-        rb_define_method(rb_klass, "set_input", RUBY_METHOD_FUNC(rb_set_input), -1);
+        rb_define_method(rb_klass, "input=", RUBY_METHOD_FUNC(rb_set_input), -1);
+        rb_define_alias(rb_klass, "set_input", "input=");
+
         rb_define_method(rb_klass, "forward", RUBY_METHOD_FUNC(rb_forward), -1);
         rb_define_method(rb_klass, "empty?", RUBY_METHOD_FUNC(rb_empty), 0);
+
+        rb_define_method(rb_klass, "get_layer_names", RUBY_METHOD_FUNC(rb_get_layer_names), 0);
       }
     }
   }
diff --git a/ext/opencv/opencv.hpp b/ext/opencv/opencv.hpp
index 7df1053..4660aa9 100644
--- a/ext/opencv/opencv.hpp
+++ b/ext/opencv/opencv.hpp
@@ -15,6 +15,10 @@ namespace rubyopencv {
     return NIL_P(value) ? (default_value) : NUM2DBL(value);
   }
 
+  inline double RTEST_DEFAULT(VALUE value, bool default_value) {
+    return NIL_P(value) ? (default_value) : (RTEST(value) ? true : false);
+  }
+
   inline VALUE HASH_LOOKUP(VALUE hash, std::string key) {
     return rb_hash_lookup(hash, ID2SYM(rb_intern(key.c_str())));
   }