mirror of
https://github.com/opencv/opencv.git
synced 2025-01-18 22:44:02 +08:00
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
This commit is contained in:
parent
51d54ad4f0
commit
629d47fcd8
@ -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()
|
||||
|
Loading…
Reference in New Issue
Block a user