From 629d47fcd85f7a6b63d0d519b08acfcfb653a0ec Mon Sep 17 00:00:00 2001
From: Lubov Batanina <piccione-mail@yandex.ru>
Date: Fri, 6 Dec 2019 21:29:57 +0300
Subject: [PATCH] Merge pull request #15988 from l-bat:custom_layer

Test create custom layer in python

* check is contiguos

* Add custom layer test

* Fix test

* Remove assert

* Move assert to pyopencv dnn

* remove assert

* Add unregister

* Fix python2

* proto to bytearray

* Fix data type
---
 modules/dnn/misc/python/test/test_dnn.py | 68 ++++++++++++++++++++++++
 1 file changed, 68 insertions(+)

diff --git a/modules/dnn/misc/python/test/test_dnn.py b/modules/dnn/misc/python/test/test_dnn.py
index 8346297c39..24a23d6ace 100644
--- a/modules/dnn/misc/python/test/test_dnn.py
+++ b/modules/dnn/misc/python/test/test_dnn.py
@@ -226,6 +226,74 @@ class dnn_test(NewOpenCVTests):
                 self.assertTrue(ret)
                 normAssert(self, refs[i], result, 'Index: %d' % i, 1e-10)
 
+    def test_custom_layer(self):
+        class CropLayer(object):
+            def __init__(self, params, blobs):
+                self.xstart = 0
+                self.xend = 0
+                self.ystart = 0
+                self.yend = 0
+            # Our layer receives two inputs. We need to crop the first input blob
+            # to match a shape of the second one (keeping batch size and number of channels)
+            def getMemoryShapes(self, inputs):
+                inputShape, targetShape = inputs[0], inputs[1]
+                batchSize, numChannels = inputShape[0], inputShape[1]
+                height, width = targetShape[2], targetShape[3]
+                self.ystart = (inputShape[2] - targetShape[2]) // 2
+                self.xstart = (inputShape[3] - targetShape[3]) // 2
+                self.yend = self.ystart + height
+                self.xend = self.xstart + width
+                return [[batchSize, numChannels, height, width]]
+            def forward(self, inputs):
+                return [inputs[0][:,:,self.ystart:self.yend,self.xstart:self.xend]]
+
+        cv.dnn_registerLayer('CropCaffe', CropLayer)
+        proto = '''
+        name: "TestCrop"
+        input: "input"
+        input_shape
+        {
+            dim: 1
+            dim: 2
+            dim: 5
+            dim: 5
+        }
+        input: "roi"
+        input_shape
+        {
+            dim: 1
+            dim: 2
+            dim: 3
+            dim: 3
+        }
+        layer {
+          name: "Crop"
+          type: "CropCaffe"
+          bottom: "input"
+          bottom: "roi"
+          top: "Crop"
+        }'''
+
+        net = cv.dnn.readNetFromCaffe(bytearray(proto.encode()))
+        for backend, target in self.dnnBackendsAndTargets:
+            if backend != cv.dnn.DNN_BACKEND_OPENCV:
+                continue
+
+            printParams(backend, target)
+
+            net.setPreferableBackend(backend)
+            net.setPreferableTarget(target)
+            src_shape = [1, 2, 5, 5]
+            dst_shape = [1, 2, 3, 3]
+            inp = np.arange(0, np.prod(src_shape), dtype=np.float32).reshape(src_shape)
+            roi = np.empty(dst_shape, dtype=np.float32)
+            net.setInput(inp, "input")
+            net.setInput(roi, "roi")
+            out = net.forward()
+            ref = inp[:, :, 1:4, 1:4]
+            normAssert(self, out, ref)
+
+        cv.dnn_unregisterLayer('CropCaffe')
 
 if __name__ == '__main__':
     NewOpenCVTests.bootstrap()