提交 629d47fc 编写于 作者: L Lubov Batanina 提交者: Alexander Alekhin

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
上级 51d54ad4
......@@ -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()
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册