From bd87eb6e66d4db4f3d994f02e1e58e9c192e824a Mon Sep 17 00:00:00 2001
From: Dmitry Kurtaev <dmitry.kurtaev+github@gmail.com>
Date: Thu, 14 Jun 2018 15:22:08 +0300
Subject: [PATCH] Import average pooling and softmax layers from Darknet

---
 modules/dnn/src/darknet/darknet_io.cpp     | 49 ++++++++++++++++++++++
 modules/dnn/test/test_darknet_importer.cpp |  5 +++
 2 files changed, 54 insertions(+)

diff --git a/modules/dnn/src/darknet/darknet_io.cpp b/modules/dnn/src/darknet/darknet_io.cpp
index 91ebb0fa8b..03805dd364 100644
--- a/modules/dnn/src/darknet/darknet_io.cpp
+++ b/modules/dnn/src/darknet/darknet_io.cpp
@@ -212,6 +212,44 @@ namespace cv {
                     fused_layer_names.push_back(last_layer);
                 }
 
+                void setAvgpool()
+                {
+                    cv::dnn::LayerParams avgpool_param;
+                    avgpool_param.set<cv::String>("pool", "ave");
+                    avgpool_param.set<bool>("global_pooling", true);
+                    avgpool_param.name = "Pooling-name";
+                    avgpool_param.type = "Pooling";
+                    darknet::LayerParameter lp;
+
+                    std::string layer_name = cv::format("avgpool_%d", layer_id);
+                    lp.layer_name = layer_name;
+                    lp.layer_type = avgpool_param.type;
+                    lp.layerParams = avgpool_param;
+                    lp.bottom_indexes.push_back(last_layer);
+                    last_layer = layer_name;
+                    net->layers.push_back(lp);
+                    layer_id++;
+                    fused_layer_names.push_back(last_layer);
+                }
+
+                void setSoftmax()
+                {
+                    cv::dnn::LayerParams softmax_param;
+                    softmax_param.name = "Softmax-name";
+                    softmax_param.type = "Softmax";
+                    darknet::LayerParameter lp;
+
+                    std::string layer_name = cv::format("softmax_%d", layer_id);
+                    lp.layer_name = layer_name;
+                    lp.layer_type = softmax_param.type;
+                    lp.layerParams = softmax_param;
+                    lp.bottom_indexes.push_back(last_layer);
+                    last_layer = layer_name;
+                    net->layers.push_back(lp);
+                    layer_id++;
+                    fused_layer_names.push_back(last_layer);
+                }
+
                 void setConcat(int number_of_inputs, int *input_indexes)
                 {
                     cv::dnn::LayerParams concat_param;
@@ -541,6 +579,17 @@ namespace cv {
                         int pad = getParam<int>(layer_params, "pad", 0);
                         setParams.setMaxpool(kernel_size, pad, stride);
                     }
+                    else if (layer_type == "avgpool")
+                    {
+                        setParams.setAvgpool();
+                    }
+                    else if (layer_type == "softmax")
+                    {
+                        int groups = getParam<int>(layer_params, "groups", 1);
+                        if (groups != 1)
+                            CV_Error(Error::StsNotImplemented, "Softmax from Darknet with groups != 1");
+                        setParams.setSoftmax();
+                    }
                     else if (layer_type == "route")
                     {
                         std::string bottom_layers = getParam<std::string>(layer_params, "layers", "");
diff --git a/modules/dnn/test/test_darknet_importer.cpp b/modules/dnn/test/test_darknet_importer.cpp
index aaa7ef5c11..e28d9dc11e 100644
--- a/modules/dnn/test/test_darknet_importer.cpp
+++ b/modules/dnn/test/test_darknet_importer.cpp
@@ -228,4 +228,9 @@ TEST(Test_Darknet, upsample)
     testDarknetLayer("upsample");
 }
 
+TEST(Test_Darknet, avgpool_softmax)
+{
+    testDarknetLayer("avgpool_softmax");
+}
+
 }} // namespace