Merge pull request justadudewhohacks#472 from bookjan/master

add new DNN Net bindings and improve examples
ledinhphuong · Dec 9, 2018 · 5b43290 · 5b43290
2 parents 559d2a6 + d0a457c
commit 5b43290
Show file tree

Hide file tree

Showing 8 changed files with 133 additions and 46 deletions.
diff --git a/cc/modules/dnn/Net.cc b/cc/modules/dnn/Net.cc
@@ -20,6 +20,10 @@ NAN_MODULE_INIT(Net::Init) {
   Nan::SetPrototypeMethod(ctor, "setInputAsync", SetInputAsync);
   Nan::SetPrototypeMethod(ctor, "forward", Forward);
   Nan::SetPrototypeMethod(ctor, "forwardAsync", ForwardAsync);
+  Nan::SetPrototypeMethod(ctor, "getLayerNames", GetLayerNames);
+  Nan::SetPrototypeMethod(ctor, "getLayerNamesAsync", GetLayerNamesAsync);
+  Nan::SetPrototypeMethod(ctor, "getUnconnectedOutLayers", GetUnconnectedOutLayers);
+  Nan::SetPrototypeMethod(ctor, "getUnconnectedOutLayersAsync", GetUnconnectedOutLayersAsync);
 
   target->Set(Nan::New("Net").ToLocalChecked(), ctor->GetFunction());
 };
@@ -63,4 +67,32 @@ NAN_METHOD(Net::ForwardAsync) {
   );
 }
 
+NAN_METHOD(Net::GetLayerNames) {
+  FF::SyncBinding(
+      std::make_shared<NetBindings::GetLayerNamesWorker>(Net::Converter::unwrap(info.This())),
+      "Net::GetLayerNames",
+      info);
+}
+
+NAN_METHOD(Net::GetLayerNamesAsync) {
+  FF::AsyncBinding(
+      std::make_shared<NetBindings::GetLayerNamesWorker>(Net::Converter::unwrap(info.This())),
+      "Net::GetLayerNamesAsync",
+      info);
+}
+
+NAN_METHOD(Net::GetUnconnectedOutLayers) {
+  FF::SyncBinding(
+      std::make_shared<NetBindings::GetUnconnectedOutLayersWorker>(Net::Converter::unwrap(info.This())),
+      "Net::GetUnconnectedOutLayers",
+      info);
+}
+
+NAN_METHOD(Net::GetUnconnectedOutLayersAsync) {
+  FF::AsyncBinding(
+      std::make_shared<NetBindings::GetUnconnectedOutLayersWorker>(Net::Converter::unwrap(info.This())),
+      "Net::GetUnconnectedOutLayersAsync",
+      info);
+}
+
 #endif
diff --git a/cc/modules/dnn/Net.h b/cc/modules/dnn/Net.h
@@ -28,6 +28,10 @@ class Net : public Nan::ObjectWrap {
 	static NAN_METHOD(SetInputAsync);
 	static NAN_METHOD(Forward);
 	static NAN_METHOD(ForwardAsync);
+  static NAN_METHOD(GetLayerNames);
+  static NAN_METHOD(GetLayerNamesAsync);
+  static NAN_METHOD(GetUnconnectedOutLayers);
+  static NAN_METHOD(GetUnconnectedOutLayersAsync);
 };
 
 #endif
diff --git a/cc/modules/dnn/NetBindings.h b/cc/modules/dnn/NetBindings.h
@@ -79,6 +79,47 @@ namespace NetBindings {
     }
   };
 
+  struct GetLayerNamesWorker : public CatchCvExceptionWorker {
+  public:
+    cv::dnn::Net self;
+    GetLayerNamesWorker(cv::dnn::Net self) {
+      this->self = self;
+    }
+
+    std::vector<std::string> returnValue;
+
+    std::string executeCatchCvExceptionWorker() {
+      std::vector<cv::String> layerNames = self.getLayerNames();
+      std::vector<std::string> strings(
+        layerNames.begin(),
+        layerNames.end());
+      returnValue = strings;
+      return "";
+    }
+
+    v8::Local<v8::Value> getReturnValue() {
+      return StringArrayConverter::wrap(returnValue);
+    }
+  };
+
+  struct GetUnconnectedOutLayersWorker : public CatchCvExceptionWorker {
+  public:
+    cv::dnn::Net self;
+    GetUnconnectedOutLayersWorker(cv::dnn::Net self) {
+      this->self = self;
+    }
+
+    std::vector<int> layerIndexes;
+
+    std::string executeCatchCvExceptionWorker() {
+      layerIndexes = self.getUnconnectedOutLayers();
+      return "";
+    }
+
+    v8::Local<v8::Value> getReturnValue() {
+      return IntArrayConverter::wrap(layerIndexes);
+    }
+  };
 
 }
 

diff --git a/cc/modules/dnn/dnn.cc b/cc/modules/dnn/dnn.cc
@@ -20,8 +20,8 @@ NAN_MODULE_INIT(Dnn::Init) {
   Nan::SetMethod(target, "blobFromImages", BlobFromImages);
   Nan::SetMethod(target, "blobFromImagesAsync", BlobFromImagesAsync);
 #if CV_VERSION_MINOR > 3
-  Nan::SetMethod(target, "readNetFromDarknet", readNetFromDarknet);
-  Nan::SetMethod(target, "readNetFromDarknetAsync", readNetFromDarknetAsync);
+  Nan::SetMethod(target, "readNetFromDarknet", ReadNetFromDarknet);
+  Nan::SetMethod(target, "readNetFromDarknetAsync", ReadNetFromDarknetAsync);
   Nan::SetMethod(target, "NMSBoxes", NMSBoxes);
 #endif
 };
@@ -92,19 +92,17 @@ NAN_METHOD(Dnn::BlobFromImagesAsync) {
 }
 
 #if CV_VERSION_MINOR > 3
-NAN_METHOD(Dnn::readNetFromDarknet)
-{
+NAN_METHOD(Dnn::ReadNetFromDarknet) {
   FF::SyncBinding(
-      std::make_shared<DnnBindings::readNetFromDarknetWorker>(),
-      "readNetFromDarknet",
+      std::make_shared<DnnBindings::ReadNetFromDarknetWorker>(),
+      "ReadNetFromDarknet",
       info);
 }
 
-NAN_METHOD(Dnn::readNetFromDarknetAsync)
-{
+NAN_METHOD(Dnn::ReadNetFromDarknetAsync) {
   FF::AsyncBinding(
-      std::make_shared<DnnBindings::readNetFromDarknetWorker>(),
-      "readNetFromDarknetAsync",
+      std::make_shared<DnnBindings::ReadNetFromDarknetWorker>(),
+      "ReadNetFromDarknetAsync",
       info);
 }
 

diff --git a/cc/modules/dnn/dnn.h b/cc/modules/dnn/dnn.h
@@ -20,8 +20,8 @@ class Dnn {
   static NAN_METHOD(BlobFromImages);
   static NAN_METHOD(BlobFromImagesAsync);
 #if CV_VERSION_MINOR > 3
-  static NAN_METHOD(readNetFromDarknet);
-  static NAN_METHOD(readNetFromDarknetAsync);
+  static NAN_METHOD(ReadNetFromDarknet);
+  static NAN_METHOD(ReadNetFromDarknetAsync);
   static NAN_METHOD(NMSBoxes);
 #endif
 };

diff --git a/cc/modules/dnn/dnnBindings.h b/cc/modules/dnn/dnnBindings.h
@@ -6,7 +6,7 @@
 namespace DnnBindings {
 
 #if CV_VERSION_MINOR > 3
-  struct readNetFromDarknetWorker : public CatchCvExceptionWorker{
+  struct ReadNetFromDarknetWorker : public CatchCvExceptionWorker{
   public:
     std::string cfgFile;
     std::string darknetModelFile = "";

diff --git a/examples/dnnDarknetYOLORealTimeObjectDetection.js b/examples/dnnDarknetYOLORealTimeObjectDetection.js
@@ -4,7 +4,7 @@
  */
 const fs = require("fs");
 const path = require("path");
-const { cv, drawBlueRect, runVideoDetection } = require("./utils");
+const { cv, runVideoDetection } = require("./utils");
 
 if (!cv.xmodules.dnn) {
   throw new Error("exiting: opencv4nodejs compiled without dnn module");
@@ -41,6 +41,13 @@ const labels = fs
 
 // initialize tensorflow darknet model from modelFile
 const net = cv.readNetFromDarknet(cfgFile, weightsFile);
+const allLayerNames = net.getLayerNames();
+const unconnectedOutLayers = net.getUnconnectedOutLayers();
+
+// determine only the *output* layer names that we need from YOLO
+const layerNames = unconnectedOutLayers.map(layerIndex => {
+  return allLayerNames[layerIndex - 1];
+});
 
 const classifyImg = img => {
   // object detection model works with 416 x 416 images
@@ -49,12 +56,9 @@ const classifyImg = img => {
   const [imgHeight, imgWidth] = img.sizes;
 
   // network accepts blobs as input
-  const inputBlob = cv.blobFromImage(img, 1 / 255.0, size, vec3, true, true);
+  const inputBlob = cv.blobFromImage(img, 1 / 255.0, size, vec3, true, false);
   net.setInput(inputBlob);
 
-  // specify two layers "yolo_16" and "yolo_23"
-  const layerNames = ["yolo_16", "yolo_23"];
-
   console.time("net.forward");
   // forward pass input through entire network
   const layerOutputs = net.forward(layerNames);
@@ -95,20 +99,25 @@ const classifyImg = img => {
 
         indices.forEach(i => {
           const rect = boxes[i];
-          const imgRect = new cv.Rect(rect.x, rect.y, rect.width, rect.height);
-          drawBlueRect(img, imgRect);
+
+          const pt1 = new cv.Point(rect.x, rect.y);
+          const pt2 = new cv.Point(rect.x + rect.width, rect.y + rect.height);
+          const rectColor = new cv.Vec(255, 0, 0);
+          const rectThickness = 2;
+          const rectLineType = cv.LINE_8;
+
+          // draw the rect for the object
+          img.drawRectangle(pt1, pt2, rectColor, rectThickness, rectLineType);
+
           const text = labels[classIDs[i]];
-          img.putText(
-            text,
-            new cv.Point(rect.x, rect.y + 0.1 * imgHeight),
-            cv.FONT_ITALIC,
-            2,
-            {
-              color: new cv.Vec(255, 0, 0),
-              thickness: 2
-            }
-          );
-          drawBlueRect(img, imgRect);
+          const org = new cv.Point(rect.x, rect.y + 15);
+          const fontFace = cv.FONT_HERSHEY_SIMPLEX;
+          const fontScale = 0.5;
+          const textColor = new cv.Vec(123, 123, 255);
+          const thickness = 2;
+
+          // put text on the object
+          img.putText(text, org, fontFace, fontScale, textColor, thickness);
         });
       }
     });

diff --git a/examples/dnnTensorflowObjectDetection.js b/examples/dnnTensorflowObjectDetection.js
@@ -5,7 +5,7 @@
 const fs = require("fs");
 const path = require("path");
 const classNames = require("./dnnTensorflowObjectDetectionClassNames");
-const { cv, drawBlueRect, runVideoDetection } = require("./utils");
+const { cv, runVideoDetection } = require("./utils");
 
 if (!cv.xmodules.dnn) {
   throw new Error("exiting: opencv4nodejs compiled without dnn module");
@@ -51,7 +51,7 @@ const classifyImg = img => {
 
   // get height and width from the image
   const [imgHeight, imgWidth] = img.sizes;
-  const numRows = outputBlob.sizes.slice(2,3);
+  const numRows = outputBlob.sizes.slice(2, 3);
 
   for (let y = 0; y < numRows; y += 1) {
     const confidence = outputBlob.at([0, 0, y, 2]);
@@ -62,22 +62,25 @@ const classifyImg = img => {
       const boxY = imgHeight * outputBlob.at([0, 0, y, 4]);
       const boxWidht = imgWidth * outputBlob.at([0, 0, y, 5]);
       const boxHeight = imgHeight * outputBlob.at([0, 0, y, 6]);
-      const imgRect = new cv.Rect(boxX, boxY, boxWidht, boxHeight);
 
-      // draw the blue rect for the object
-      drawBlueRect(img, imgRect);
+      const pt1 = new cv.Point(boxX, boxY);
+      const pt2 = new cv.Point(boxWidht, boxHeight);
+      const rectColor = new cv.Vec(23, 230, 210);
+      const rectThickness = 2;
+      const rectLineType = cv.LINE_8;
+
+      // draw the rect for the object
+      img.drawRectangle(pt1, pt2, rectColor, rectThickness, rectLineType);
+
+      const text = `${className} ${confidence.toFixed(5)}`;
+      const org = new cv.Point(boxX, boxY + 15);
+      const fontFace = cv.FONT_HERSHEY_SIMPLEX;
+      const fontScale = 0.5;
+      const textColor = new cv.Vec(255, 0, 0);
+      const thickness = 2;
 
       // put text on the object
-      img.putText(
-        className,
-        new cv.Point(boxX, boxY + 0.1 * imgHeight),
-        cv.FONT_ITALIC,
-        2,
-        {
-          color: new cv.Vec(255, 0, 0),
-          thickness: 2
-        }
-      );
+      img.putText(text, org, fontFace, fontScale, textColor, thickness);
     }
   }